Save the optimizer state dict together with the model state dict

[DanielCoelho112]

According to https://pytorch.org/tutorials/beginner/saving_loading_models.html, when saving the checkpoint we should also save the state dict of the optimizer.