[FIXED] Save a model on gpu and load it on CPU

Issue

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:

Solution

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))

Answered By - Alex