added saving tests to cpu

This commit is contained in:
William Falcon
2019-07-26 12:14:58 -04:00
parent 2ee8f157ce
commit 1a835969a6
2 changed files with 27 additions and 16 deletions
@@ -86,7 +86,7 @@ class TrainerIO(object):
# --------------------
# HPC IO
# --------------------
def enable_auto_hpc_walltime_manager(self): # pragma: no cover
def enable_auto_hpc_walltime_manager(self):
if self.cluster is None:
return
@@ -157,6 +157,8 @@ class TrainerIO(object):
# do the actual save
torch.save(checkpoint_dict, filepath)
return filepath
def hpc_load(self, folderpath, on_gpu):
filepath = '{}/hpc_ckpt_{}.ckpt'.format(folderpath, self.max_ckpt_in_folder(folderpath))