Trains do patch the torch save function 🙂

If you like, you can save the model for each epoch by having a unique name for it. The model will be saved in the output_uri path you have in the Task.init command.
For example, this code will save a model for every epoch:

for epoch in range(num_of_epoch): # Create a model torch.save(model, "model_epoch_{}".format(epoch))