Epoch is not saved to the Torch checkpoint

[martinlyra]

Whenever predict_ros.py is run with our own trained data:

Loading ckpt from  /home/marcusmartin/artifacts/0/model_best_val.pth.tar
Traceback (most recent call last):
  File "predict_ros.py", line 109, in <module>
    tracker = Tracker(dataset_info, images_mean, images_std,ckpt_dir,trans_normalizer=dataset_info['max_translation'],rot_normalizer=dataset_info['max_rotation'])
  File "/home/marcusmartin/iros20-tracking/predict.py", line 152, in __init__
    print('pose track ckpt epoch={}'.format(checkpoint['epoch']))
KeyError: 'epoch'

The culprit seems to be

iros20-6d-pose-tracking/problems.py

Line 142 in 1b63d0d

checkpoint_data = {'state_dict': self.model.state_dict()}

and

iros20-6d-pose-tracking/problems.py

Line 149 in 1b63d0d

checkpoint_data = {'state_dict': self.model.state_dict()}

Our torch version when we trained the data was 2.0.1 - if you were using an older version, may I suggest you to add a requirements.txt too?

[wenbowen123]

can you pull the code and retry? This issue should not be related to pytorch version. Are you running inside the docker?

[martinlyra]

Nope, we run natively on a computer that has access to the robot ROS. - The fix seems to work fine!