How to load a saved model defined in a function in PyTorch in Colab?
Here is a sample code of my training function(unnecessary parts are deleted):
I am trying to save my model data_gen in the torch.save(), and after running the train_dmc function, I can find the checkpoint file in the directory.
def train_dmc(loader,loss):


 
  data_gen = DataGenerator().to(device)

  data_gen_optimizer = optim.Rprop(para_list, lr=lrate)


  savepath='/content/drive/MyDrive/'+loss+'checkpoint.t7'
  state = {
            'epoch': epoch,
            'model_state_dict': data_gen.state_dict(),
            'optimizer_state_dict': data_gen_optimizer.state_dict(),
            'data loss': data_loss,
            'latent_loss':latent_loss
            }
  torch.save(state,savepath)

My question is that how to load the checkpoint file to continue training if Google Colab disconnects.
Should I load data_gen or train_dmc(), it is my first time using this and I am really confused because the data_gen is defined inside another function. Hope someone can help me with explanation
data_gen.load_state_dict(torch.load(PATH))
data_gen.eval()

#or

train_dmc.load_state_dict(torch.load(PATH))
train_dmc.eval()


As the state variable is a dictionary, So try saving it as:
with open('/content/checkpoint.t7', 'wb') as handle:
    pickle.dump(state, handle, protocol=pickle.HIGHEST_PROTOCOL)

Initiate your model class as data_gen = DataGenerator().to(device).
And load the checkpoint file as:
import pickle
file = open('/content/checkpoint.t7', 'rb')
loaded_state = pickle.load(file)

Then you can load the state_dict using data_gen = loaded_state['model_state_dict']. This will load the state_dict to the model class!
