@cpeters Great, thank you so much!
I am having difficulties applying it to my scenario. I only have two epochs but each contains 146,447 batches. So, I would need to do it on a batch level rather than epoch level. I have a train() and evaluate() function that are called like this
for epoch in range(params['epochs']):
     
      print('\n Epoch {:} / {:}'.format(epoch + 1, params['epochs']))
      
      #train model
      train_loss = train(scheduler, optimizer)
      
      #evaluate model
      valid_loss = evaluate()
      
      #save the best model -> OLD
      #this only saves the model at the end of each epoch. I want to change it.
      '''if valid_loss < best_valid_loss:
          best_valid_loss = valid_loss
          torch.save(model.state_dict(), model_file)'''
      
      # append training and validation loss
      train_losses.append(train_loss)
      valid_losses.append(valid_loss)
      
      print(f'\nTraining Loss: {train_loss:.3f}')
      print(f'Validation Loss: {valid_loss:.3f}')

I could create the checkpoint within train() but then won’t have access to valid_loss within train().
Any idea how I could resolve this? 
I used to create the checkpoint within train() like this:
# checkpoint
    if step % 300 == 0 and not step == 0:
      checkpoint = {'state_dict': model.state_dict(), 
                      'optimizer': optimizer.state_dict(), 
                      'train_loss': step_loss,
                      #'val_Loss': valid_loss, 
                      }
      save_checkpoint(checkpoint)

It saves the stated params all 300 steps. This works but if I load from it, it starts from the beginning of the dataset again. So, I would need to create something like range(start_batch, end_batch)? Also, I do not have access to the validation loss within train() as mentioned before.