How do you save and load model if Autoencoder and Classifier is there?

[Praveenk8051]

Hello Abhishek,

Your posts have been really helpful. Can you please let me know how do you load the saved model if 2 stages are involved. Meaning how do load the state_dict

[abhisheksambyal]

Assuming you have saved your model/state like this:

torch.save(
        {
            'epoch': Config.num_epochs,
            'model_state_dict': modelCAN.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            'global_trnloss': global_trn_loss,
            'global_valloss': global_val_loss
        }, MODEL_SAVE_PATH)

You can use this to load your saved model:

checkpoint = torch.load('./model_colab_300.pt', map_location='cuda')
model.load_state_dict(checkpoint['model_state_dict'])