title: How to load a saved model defined in a function in PyTorch in Colab?
body: <p>Here is a sample code of my training function(unnecessary parts are deleted):</p>
  <p>I am trying to save my model <strong>data_gen</strong> in the torch.save(), and after running the train_dmc function, I can find the checkpoint file in the directory.</p>
  <pre><code>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)
  </code></pre>
  <p>My question is that how to load the checkpoint file to continue training if Google Colab disconnects.</p>
  <p>Should I load <strong>data_gen</strong> or <strong>train_dmc()</strong>, 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</p>
  <pre><code>data_gen.load_state_dict(torch.load(PATH))
  data_gen.eval()
  
  #or
  
  train_dmc.load_state_dict(torch.load(PATH))
  train_dmc.eval()
  </code></pre>