Issue with different devices

[obakumenkofh]

Dear Authors,

I am using your repository for my re-implementation of the CLAM model. You can see the citation in the environment file here: https://github.com/mahmoodlab/CLAM/blob/master/env.yml.

I would like to ask a question regarding support for different devices in PyTorch. I encountered an issue where, if ( x ) and ( y ) are located on the "cuda" device, the SmoothTopkSVM cannot be computed because the "labels" variable is on the "cpu". I reused the explanatory code from a previous issue and added a device line.

import torch
from topk.svm import SmoothTopkSVM

# hyper-parameters
topk = 10  # top-10 classification
alpha = 1  # margin parameter of hinge loss, 1 is a default value
tau = 0.1  # smoothing temperature parameter (see paper for details)

# define the device
if torch.cuda.is_available():
    device = torch.device('cuda')

# random data
batch_size = 3
n_dim = 20
n_classes = 100
x = torch.randn(batch_size, n_dim).to(device)
y = torch.randint(0, n_classes, size=(batch_size,)).long().to(device)

# model
model = torch.nn.Linear(n_dim, n_classes).to(device)

# loss function
loss_fn = SmoothTopkSVM(n_classes, alpha=alpha, tau=tau, k=topk).to(device)

# forward pass
loss = loss_fn(model(x), y)

# backward pass
loss.backward()

# print gradient
print(model.weight.grad)

If I run this, the error RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! comes from this line, where the "labels" variable is not on "cuda":
https://github.com/oval-group/smooth-topk/blob/master/topk/utils.py#L23.

I could fix it with:

labels = ag.Variable(labels, requires_grad=False).to(x.device)

and also here: https://github.com/oval-group/smooth-topk/blob/master/topk/functional.py#L35 with:

x = x + delta(y, labels.to(x.device), alpha) - x.gather(1, y[:, None])

Thank you for your time and help.
Best regards :)