Epochs to be stored in different files

Hi Folks,
How would I be able to split my epochs to different files?
For eg-
If I am training my model for 20 epochs, I would like that at every five epochs, they get stored in 4 different files (.pt).
How should I approach this?

Below is a code I am trying, any feedback would be appreciated-
model = nn.Sequential(

nn.Linear(1, 30),

nn.Tanh(),

nn.Linear(30, 30),

nn.Tanh(),

nn.Linear(30, 1),

)

epochs=20
#Simple network is a function I have defined to train my model, it works fine.
results_pd = simple_network(model, loss_func, train_set, epochs=20,

test_set, checkpoint_file=‘model.pt’,

score_funcs={‘Acc’:accuracy_score})

The point in the code where you are saving it could go something like this,

for e in range(epochs):
      .
      .
      .
      if e % 5 == 0:
         save_file_name = "{}_state.pt".format(e+1)
         torch.save(model_state_dict, save_file_name)
     .
     .
1 Like

Thanks a lot, it helped!