Change training checkpoint's format #51
Description
At present, we are saving callbacks_state, epoch_state, model_state, and training_state in the .pt format. However, it would be better to create a separate folder for each epoch and save the model_state in .pt format and the other states in .json format with the folder name being epoch_<epoch_idx>.