Unable to load model in Pytorch in another notebook

#1

I have trained M5 model for the purposes of sentence classification. As I have saved it using torch.save(PATH) , I am no more able to load it in another notebook using torch.load(PATH) . The error, that I see is following:

AttributeError: Can't get attribute 'M5' on <module '__main__'>
#2

You should use

torch.save(model.state_dict(), PATH)

and

model.load_state_dict(torch.load(PATH))
#3

Thanks! Thats working.