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

Saeid_Sadeghnejad: <p>Hi everyone,</p>
<p>I want to load my previous saved model and continue training. This is the code for saving the model. The network is being trained on GPU (and works well there).</p>
<pre><code class="lang-auto"># Save model
torch.save({
       'epoch': epoch,
       'model_state_dict': model.state_dict(),
       'optimizer_state_dict': optimizer.state_dict(),
       'loss': loss, }, f'{chkp_Path}/epoch{epoch}_{datetime.now().strftime("%Y%m%d-%H%M%S")}.pt') 
</code></pre>
<p>During the loading phase, I load the model on CPU (to prevent out of memory issue, <a href="https://discuss.pytorch.org/t/cuda-error-out-of-memory-when-load-models/38011">like</a>) and then transfer it back to GPU:</p>
<pre><code class="lang-auto">chkPath = 'Path_to_chk_file.pt'
checkpoint = torch.load(chkPath, map_location='cpu')
start_epoch = checkpoint['epoch']
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
losslogger = checkpoint['loss']

device = torch.device(str("cuda:0") if torch.cuda.is_available() else "cpu") 
model = model.to(device) # the network should be on GPU, next(model.parameters()).is_cuda==False
</code></pre>
<p>Then , when I start the training phase, I receive the following error. Is there any way to find out where those torch tensors on CPU come from? BTW, the data loader is on GPU during the training.</p>
<p><strong>RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!</strong></p>
<p>The above error happen at this line and after loss calculaiton</p>
<pre><code class="lang-auto">opt.step()
</code></pre>

ptrblck: <p>These lines of code:</p>
<pre><code class="lang-python">device = torch.device(str("cuda:0") if torch.cuda.is_available() else "cpu") 
model = model.to(device) # the network should be on GPU, next(model.parameters()).is_cuda==False
</code></pre>
<p>look interesting, as it seems that parameters are not passed to the GPU.<br>
Could you make sure <code>device</code> is indeed <code>cuda:0</code> (I assume so)?<br>
If so, could you post the model definition, please?</p>

Saeid_Sadeghnejad: <aside class="quote no-group" data-username="ptrblck" data-post="3" data-topic="143897">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>ition, please</p>
</blockquote>
</aside>
<p>Thanks for your reply. Device is “cuda:0”. I have explicitly set it to 'cuda:0" and also checked the GPU load. It is running on GPU.<br>
Please find the model. It is a (Vibrational) Auto Encoder for 3D images (adopted from <a href="https://avandekleut.github.io/vae/" rel="noopener nofollow ugc">here</a>).</p>
<pre><code class="lang-auto">class Encoder(nn.Module):
    def __init__(self, latent_dims, VE):
        super(Encoder, self).__init__()
        self.CNN = nn.Sequential(
                        nn.Conv3d(1, 64, kernel_size=5, padding=2, stride=1),
                        nn.ReLU(),
                        nn.MaxPool3d(2),
                       
                        nn.Conv3d(64, 128, kernel_size=5, padding=2, stride=1),
                        nn.ReLU(),
                        nn.MaxPool3d(2),
                                    
                        nn.Conv3d(128, 256, kernel_size=5, padding=2, stride=1),
                        nn.ReLU(),
                        nn.MaxPool3d(2),
                            
                        nn.Conv3d(256, 256, kernel_size=5, padding=2, stride=1),
                        nn.LeakyReLU(0.1),
                        nn.MaxPool3d(2),
                       
                        nn.Flatten(),
            )
        
        self.linear1 = nn.Linear(8**3*256, latent_dims)
        self.linear2 = nn.Linear(8**3*256, latent_dims)

        self.N = torch.distributions.Normal(0, 1)
       
        self.kl = 0
        self.VE = VE
        
    def forward(self, x):
        if self.VE:
            x = self.CNN(x)
            mu =  self.linear1(x)
            sigma = torch.exp(self.linear2(x)) # exp: to ensure the result is positive
            z = mu + sigma*self.N.sample(mu.shape)
            self.kl = (sigma**2 + mu**2 - torch.log(sigma) - 1/2).sum()
        else:
            z = self.CNN(x)
            z = self.linear1(z)
        return z
    
class Decoder(nn.Module):
    def __init__(self, latent_dims):
        super(Decoder, self).__init__()
        self.Linear = nn.Linear(latent_dims, 8**3*256)
        self.CNN = nn.Sequential(
                            nn.ConvTranspose3d(256, 256, kernel_size=4, stride=2, padding=1),
                            nn.ReLU(),
                           
                            nn.ConvTranspose3d(256, 128, kernel_size=4, stride=2, padding=1),
                            nn.ReLU(),
                            
                            nn.ConvTranspose3d(128, 64, kernel_size=4, stride=2, padding=1),
                            nn.ReLU(),
                            
                            nn.ConvTranspose3d(64, 1, kernel_size=4, stride=2, padding=1),
                            nn.Sigmoid()
                          )
    def forward(self, z):
        z = self.Linear(z)
        z = z.view((-1, 256, 8, 8, 8))
        z = self.CNN(z)
        return z
    
class VariationalAutoencoder(nn.Module):
    def __init__(self, latent_dims, VE):
        super(VariationalAutoencoder, self).__init__()
        self.encoder = Encoder(latent_dims, VE)
        self.decoder = Decoder(latent_dims)

    def forward(self, x):
        z = self.encoder(x)
        return self.decoder(z)
</code></pre>
<p>Tnx for your help.</p>

ptrblck: <p>Thanks for the code! Your model seems to be fine in my setup:</p>
<pre><code class="lang-python"># save    
model = VariationalAutoencoder(1, 1)
model.to('cuda:0')
print(next(model.parameters()).device)
# &gt; cuda:0

torch.save(model.state_dict(), 'tmp.pt')


# load in another script
import torch
import torch.nn as nn

model = VariationalAutoencoder(1, 1)
checkpoint = torch.load('tmp.pt', map_location='cpu')
model.load_state_dict(checkpoint)
# &gt; &lt;All keys matched successfully&gt;

model.to('cuda:0')
print(next(model.parameters()).device)
# &gt; cuda:0
</code></pre>
<p>Could you check, if you are able to run into the same issue using my code snippet and if not where the difference might be?</p>

Saeid_Sadeghnejad: <aside class="quote no-group" data-username="ptrblck" data-post="5" data-topic="143897">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>nto the same issue</p>
</blockquote>
</aside>
<p>Tnx for your prompt reply. I am doing very similar to your code. After the loading of the checkpoint, the model is on GPU (print(next(model.parameters()).device)==cuda:0). However, when the training code reaches to optimizer.step(), inside of the training function, it throws the error. It means that input and model are both at the same machine because loss can be computed for the first minibatch at least. However, something inside of the loaded optimizer should be wrong. It seems that some of the Optmizer parameters should be in CPU and some of the others in GPU simultaneously. I am not expert how pytorch optimizer is working. Is there any similar code to model.to(device) for optimizer (like opt.to(device)) available?</p>

Saeid_Sadeghnejad: <p>I googled the idea of moving optimizer parameters from <strong>cpu</strong> to <strong>cuda:0</strong> and found this bunch of code from <a href="https://github.com/pytorch/pytorch/issues/8741#issuecomment-402129385" rel="noopener nofollow ugc">here</a>:</p>
<pre><code class="lang-auto">def optimizer_to(optim, device):
    # move optimizer to device
    for param in optim.state.values():
        # Not sure there are any global tensors in the state dict
        if isinstance(param, torch.Tensor):
            param.data = param.data.to(device)
            if param._grad is not None:
                param._grad.data = param._grad.data.to(device)
        elif isinstance(param, dict):
            for subparam in param.values():
                if isinstance(subparam, torch.Tensor):
                    subparam.data = subparam.data.to(device)
                    if subparam._grad is not None:
                        subparam._grad.data = subparam._grad.data.to(device)
</code></pre>
<p>Surprisingly, it solved my problem. <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_smiling_face.png?v=12" title=":slightly_smiling_face:" class="emoji" alt=":slightly_smiling_face:"> But I do not why the optimizer parameters still remain in ‘cpu’ when the uploaded model moves from cpu to gpu? any idea?</p>
<p>This is how I solved the problem:</p>
<pre><code class="lang-auto">def load_checkpoint(model, optimizer, losslogger, device, filename='checkpoint.pt'):
    # https://discuss.pytorch.org/t/loading-a-saved-model-for-continue-training/17244/3
    start_epoch = 0
    if os.path.isfile(filename):
        print("=&gt; loading checkpoint '{}'".format(filename))
        checkpoint = torch.load(filename, map_location=device)
        start_epoch = checkpoint['epoch']
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        losslogger = checkpoint['loss']
        print("=&gt; loaded epoch: {} of checkpoint: '{}' "
                  .format(filename, checkpoint['epoch']))
    else:
        print("=&gt; no checkpoint found at '{}'".format(filename))
        exit()
    return model, optimizer, start_epoch, losslogger  

 if Phase=='continue_training':
        loss = 0
        model, optimizer, start_epoch, loss = load_checkpoint(model, optimizer, loss, 'cpu', chkPath)
        model = model.to(device)         # move back model to GPU
        assert next(model.parameters()).device == 'cuda:0'        
        optimizer_to(optimizer, device)  # *** move optimizer to GPU ***
</code></pre>

ptrblck: <p>Good to hear you’ve solved the issue!<br>
I think loading the optimizer on the CPU via <code>map_location</code> is causing the trouble.</p>

Saeid_Sadeghnejad: <p>Thank you for the discussion.<br>
That’s right. The Pytorch optimizer is not straight forward and something internally happens. Maybe there is an internal bug…</p>

guangzhi: <p>I got the same issue today, and managed to fix it by changing the order of the 3 steps to the following:</p>
<ol>
<li>call <code>model.to(device)</code> first, then</li>
<li>
<code>model.load_state_dict()</code> second.</li>
<li>
<code>opt.load_state_dict()</code> last.</li>
</ol>
<p>I loaded the <code>state_dict</code> of the optimizer as the last step, but switching 2 and 3 doesn’t seem to make a difference. I guess it is important to move the model as the 1st step. Judging from your original post this might be the issue, maybe you could give it a try.</p>
<p>Btw I also added <code>map_location="cpu"</code> to <code>torch.load()</code>.</p>

ekmungi: <p>This solution worked for me. I didn’t add,</p>
<aside class="quote no-group" data-username="guangzhi" data-post="10" data-topic="143897">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/g/e274bd/48.png" class="avatar"> guangzhi:</div>
<blockquote>
<p><code>map_location="cpu"</code> to <code>torch.load()</code>.</p>
</blockquote>
</aside>
<p>But changed the order as suggested in reply by <a class="mention" href="/u/guangzhi">@guangzhi</a> .</p>

khaidoan25: <p>The solution of <a class="mention" href="/u/guangzhi">@guangzhi</a> works for me.</p>