RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
This comes when I run the example notebook from this github repo.
Here is the code:
batch_size = 128
epochs = 1000 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True)
model = SCARF( input_dim=train_ds.shape[1], emb_dim=16, corruption_rate=0.6, ).to(device) optimizer = Adam(model.parameters(), lr=0.001) ntxent_loss = NTXent()
loss_history = []
for epoch in range(1, epochs + 1): epoch_loss = train_epoch(model, ntxent_loss, train_loader, optimizer, device, epoch) loss_history.append(epoch_loss)
RuntimeError Traceback (most recent call last) Cell In [7], line 7 4 loss_history = [] 6 for epoch in range(1, epochs + 1): ----> 7 epoch_loss = train_epoch(model, ntxent_loss, train_loader, optimizer, device, epoch) 8 loss_history.append(epoch_loss)
File ~/pytorch-scarf/example/../example/utils.py:23, in train_epoch(model, criterion, train_loader, optimizer, device, epoch) 20 emb_anchor, emb_positive = model(anchor, positive) 22 # compute loss ---> 23 loss = criterion(emb_anchor, emb_positive) 24 loss.backward() 26 # update model weights
File /opt/tljh/user/lib/python3.9/site-packages/torch/nn/modules/module.py:1130, in Module._call_impl(self, *input, **kwargs) 1126 # If we don't have any hooks, we want to skip the rest of the logic in 1127 # this function, and just call forward. 1128 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks 1129 or _global_forward_hooks or _global_forward_pre_hooks): -> 1130 return forward_call(*input, **kwargs) 1131 # Do not call functions when jit is used 1132 full_backward_hooks, non_full_backward_hooks = [], []
File ~/pytorch-scarf/example/../scarf/loss.py:39, in NTXent.forward(self, z_i, z_j) 37 mask = (~torch.eye(batch_size * 2, batch_size * 2, dtype=torch.bool)).float() 38 numerator = torch.exp(positives / self.temperature) ---> 39 denominator = mask * torch.exp(similarity / self.temperature) 41 all_losses = -torch.log(numerator / torch.sum(denominator, dim=1)) 42 loss = torch.sum(all_losses) / (2 * batch_size)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
When I run the code on a CPU only machine I don't get this same error. Because of how the data is created I am not able to confirm what tensor type it is (maybe this is the problem). I've confirmed that both emb_anchor and emb_positive before they are passed into criterion() are cuda (as suggested by this post to be a possible solution)
The problem is in scarf/loss.py file. You should replace the line:
mask = (~torch.eye(batch_size * 2, batch_size * 2, dtype=torch.bool)).float()
with
mask = (~torch.eye(batch_size * 2, batch_size * 2, dtype=torch.bool)).float().to(z_i.device)
The author forgot to move mask tensor to z_i.device