Save a model on gpu and load it on CPU

Carlos Jo 0
Viewed 20

I have a model which I train based on the torch in GPU. Now, i want to upload it on cpu. I am using this code to save an load the model.

Here is my model and the training phase:

model              = VAE(input_size ,  lead,     hidden_dim,hidden_dim1,hidden_dimd,  latent_dim, device,  num_hidden= lead).to(device)

optimizer          = torch.optim.Adam(model.parameters(), lr =learning_rate )

losses, kl_loss, l_loss, validate_loss = trainv(model, device, epochs, train_iterator, optimizer, validate_iterator) 

model = VAE()
torch.save(model.state_dict(), "model.pt")

#load
device = torch.device('cpu')
model = VAE()
model.load_state_dict(torch.load(PATH, map_location=device))

Here is the error:

TypeError: __init__() missing 8 required positional arguments:
gpu pytorch python
2 Answers

The error says TheModelClass is not defined, this means the class of your model is not available in the current file. You commonly define a class e.g. as the one for your model TheModelClass in a separate module that needs to be imported first.

So since you did not provide more information about your code, I'd just suggest searching all the files in your project for the definition of TheModelClass, and importing it from the correspdonding module into your current file.

flawr0

Did you define your model class before torch.save() ? This code works without errors in Google Collab:

import torch
import torch.nn as nn

class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        
        self.linear = nn.Linear(125, 1)

    def forward(self, src):
        
        output = self.linear(src)

        return

model = TheModelClass()

torch.save(model.state_dict(), "model.pt") # you need to define your model before
device = torch.device('cpu')
model = TheModelClass() # you even don't need to redefine your model
model.load_state_dict(torch.load('/content/model.pt', map_location=device))
userOWH110
Related