Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

prem9: <h1>Loading model from torchvision</h1>
<p>model = models.segmentation.deeplabv3_resnet50(pretrained=True, progress=False, num_classes=21, aux_loss=None)<br>
model.classifier[4] = nn.Conv2d(256, 1, (1, 1), (1, 1))<br>
model.aux_classifier[4] = nn.Conv2d(256, 1, (1, 1), (1, 1))</p>
<h1>Freeze initial layer for finetuning</h1>
<p>idx = 0<br>
for name, param in model.named_parameters():<br>
# print(idx, name)<br>
if idx &lt; 129:<br>
param.requires_grad_ = False<br>
else:<br>
break<br>
idx += 1</p>
<h1>Defining device to cuda if GPU is avaliable else to cpu</h1>
<p>device = torch.device(‘cuda’ if torch.cuda.is_available() else ‘cpu’)</p>
<h1>Defining loss criterion</h1>
<p>criterion = FocalTverskyLoss(alpha=0.9)</p>
<h1>Defining optimizer to update model params, Adam’s a good default</h1>
<p>optimizer = optim.Adam(model.parameters())</p>
<h1>Learning rate scheduler to update lr when loss stops improving</h1>
<p>scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=2)</p>
<h1>Path at which to store model and required configs</h1>
<p>model_save_path = ‘Models/’</p>
<h1>If model file exits load previous training configs (model and other stuff) to continue training else start from 1</h1>
<p>if os.path.exists(model_save_path + ‘deeplabv3_resnet50_train.pt’):<br>
checkpoint = torch.load(model_save_path + ‘deeplabv3_resnet50_train.pt’)<br>
start = checkpoint[‘epoch’] + 1<br>
min_val_loss = checkpoint[‘val_loss’]<br>
model.load_state_dict(checkpoint[‘model_state_dict’])<br>
optimizer.load_state_dict(checkpoint[‘optimizer_state_dict’])<br>
scheduler.load_state_dict(checkpoint[‘scheduler_state_dict’])<br>
else:<br>
start = 1<br>
min_val_loss = 100</p>
<p>if torch.cuda.device_count() &gt; 1:<br>
print("Let’s use ",torch.cuda.device_count(),“GPUS!”)<br>
model = nn.DataParallel(model, device_ids=[0,1]).cuda()</p>
<h1>Move model to device</h1>
<h1>model.to(device)</h1>
<h1>Total epochs to train for</h1>
<p>epochs = 100</p>
<h1>Training loop</h1>
<p>for epoch in range(start, epochs + 1):<br>
train_loss = 0.<br>
val_loss = 0.</p>
<pre><code># Switch model to training mode
model.train()

# Forward pass through dataset
for imgs, masks in train_loader:
    # Move batch data to device
    imgs, masks = imgs.to(device), masks.to(device)
    
    # Clear previous gradients
    optimizer.zero_grad()
    
    # Forward pass the batch and get predictions
    preds = model(imgs)['out']
    
    # Calculate Loss
    loss = criterion(preds, masks)
    
    # Add to calculate loss for whole dataset
    train_loss += (loss.item() * imgs.size(0))
    
    # Backpropagate gradients
    loss.backward()
    
    # Make weight updates
    optimizer.step()
    
    # Empty cuda cache to clear useless data from VRAM for better utilization
    torch.cuda.empty_cache()

# Switch model to inference mode
model.eval()

# Forward pass through validation dataset
with torch.no_grad():
    for imgs, masks in val_loader:
        imgs, masks = imgs.to(device), masks.to(device)
        
        preds = model(imgs)['out']
        
        loss = criterion(preds, masks)
        val_loss += (loss.item() * imgs.size(0))
        
        torch.cuda.empty_cache()
  
# Average loss on the dataset
train_loss /= len(train_dataset)
val_loss /= len(val_dataset)

# Change lr if loss is not improving
scheduler.step(val_loss)

# If loss is decreasing then store model and other configs in file else not    
if val_loss &lt; min_val_loss:
    min_val_loss = val_loss
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict' : scheduler.state_dict(),
        'val_loss' : min_val_loss
        }, model_save_path + 'deeplabv3_resnet50_train.pt')
    torch.save(model, model_save_path + 'deeplabv3_resnet50_infer.pt')
    
# Print epoch number and train loss
print('Epoch {}:\tTrain Loss: {}\tVal Loss: {}'.format(epoch, train_loss, val_loss))
</code></pre>
<p>It gives error Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu</p>

ptrblck: <p>Could you post the complete error message so that we could use the stack trace to narrow down the line of code, which raises the error?</p>
<p>You can post code snippets by wrapping them into three backticks ```, which makes debugging easier <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=9" title=":wink:" class="emoji" alt=":wink:"></p>

prem9: <hr>
<p>RuntimeError                              Traceback (most recent call last)<br>
 in <br>
28<br>
29         # Make weight updates<br>
—&gt; 30         optimizer.step()<br>
31<br>
32         # Empty cuda cache to clear useless data from VRAM for better utilization</p>
<p>~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)<br>
13         def decorate_context(*args, **kwargs):<br>
14             with self:<br>
—&gt; 15                 return func(*args, **kwargs)<br>
16         return decorate_context<br>
17</p>
<p>~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/optim/adam.py in step(self, closure)<br>
97<br>
98                 # Decay the first and second moment running average coefficient<br>
—&gt; 99                 exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)<br>
100                 exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)<br>
101                 if amsgrad:</p>
<p>RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!</p>

jipson7: <p>I have a very similar error. The line causing it is optimizer.step(), but it only seems to happen after loading from a checkpoint. I am sure to move the model to the proper device after loading, is there something else we need to call <code>.to</code> on?</p>

jipson7: <p>This worked for me: <a href="https://github.com/pytorch/pytorch/issues/2830#issuecomment-718816292" class="inline-onebox" rel="noopener nofollow ugc">optimizer load_state_dict() problem? · Issue #2830 · pytorch/pytorch · GitHub</a></p>

Oshrin: <p>I have a same error like you. I try to firstly exec <code>model.cuda()</code> and then exec <code>model.load_state_dict(...) optimizer.load_state_dict(...)</code>. It works for me.</p>

D_c_Anh_Dao: <p>thanks you alot. I have the same issue. It work for me.</p>

EMU1337X: <p>I love u bro, truly. works for me well.</p>

Chetan_Annam: <p>Thanks a lot, this solution works for me. But I still don’t understand why this is working <img src="https://discuss.pytorch.org/images/emoji/apple/sweat_smile.png?v=12" title=":sweat_smile:" class="emoji" alt=":sweat_smile:" loading="lazy" width="20" height="20"></p>