I am confused how to save model for the purpose of resume training
I have a function
def fit(model,train,val):
.....
.....
return loss
I configure model as follow
model=MyModel.to(device)
criterion = nn.BCEWithLogitsLoss()
opt=torch.optim.AdamW(params=model.parameters(),lr=0.001)
Then I call function as
loss=fit(model,train,val)
Do I need to return optimizer and model from my fit function to save checkpoint as follow
state = {
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': opt.state_dict(),
}
savepath='0to5.pt'
torch.save(state,savepath)