The abstract elbo.ElboModel
An ElboModel is an abstract class that allows ELBO service to automatically checkpoint your training.
Extending your model class
Extend the abstractElboModel
along with nn.Model
in your PyTorch model class. With this you will be required to implement two methods:
save_state
- This method should save the state of the model and other state information needed.load_state
- This method should load the state of the model from the input directory.
Good to know: These methods will be called by the training loop at periodic intervals to keep saving the state of training. Please make sure anything thats needed for a training to resume from a previous checkpoint is saved and loaded through this method.
Last updated