From b86d223889992cccbb4ca68537a464df88dbfef7 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Fri, 25 Oct 2019 08:57:05 -0400 Subject: [PATCH] makes checkpoint process safe (#431) --- pytorch_lightning/trainer/trainer_io.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/pytorch_lightning/trainer/trainer_io.py b/pytorch_lightning/trainer/trainer_io.py index 5b483760..b8de2bbf 100644 --- a/pytorch_lightning/trainer/trainer_io.py +++ b/pytorch_lightning/trainer/trainer_io.py @@ -137,7 +137,13 @@ class TrainerIOMixin(object): checkpoint = self.dump_checkpoint() # do the actual save - torch.save(checkpoint, filepath) + try: + torch.save(checkpoint, filepath) + except AttributeError: + if 'hparams' in checkpoint: + del checkpoint['hparams'] + + torch.save(checkpoint, filepath) def restore(self, checkpoint_path, on_gpu): @@ -283,7 +289,14 @@ class TrainerIOMixin(object): model.on_hpc_save(checkpoint) # do the actual save - torch.save(checkpoint, filepath) + # TODO: fix for anything with multiprocess DP, DDP, DDP2 + try: + torch.save(checkpoint, filepath) + except AttributeError: + if 'hparams' in checkpoint: + del checkpoint['hparams'] + + torch.save(checkpoint, filepath) return filepath