It is recommended to remove the "device" parameter from the code

[zhang090210]

class AUCMLoss(torch.nn.Module):
    def __init__(self, 
                 margin=1.0, 
                 imratio=None,
                 version='v1',
                 device=None):
        super(AUCMLoss, self).__init__()
        if not device:
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            self.device =  device

The lightning framework supports the use of multiple devices, and users who do not use the framework can use the to method to move to the desired device, and the "device" parameter makes the lightning framework use RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:2 and cuda:0!
So I suggest removing the "device" parameter

[zhang090210]

In addition, the tensors defined in the class will not be actively transferred to "device", and they should be changed to "nn. Parameter", the content comes from AI, please fully verify before adopting

example:
origin:

self.a = torch.zeros(1, dtype=torch.float32, requires_grad=True)
self.b = torch.zeros(1, dtype=torch.float32,  requires_grad=True) 
self.alpha = torch.zeros(1, dtype=torch.float32, requires_grad=True)

change

self.a = nn.Parameter(torch.zeros(1, dtype=torch.float32))
self.b = nn.Parameter(torch.zeros(1, dtype=torch.float32))
self.alpha = nn.Parameter(torch.zeros(1, dtype=torch.float32))

[optmai]

Thank you for the suggestion! We will look into this issue and improve the library.