Skip to content

Commit

Permalink
Fix general error catch with focused file check
Browse files Browse the repository at this point in the history
  • Loading branch information
tarepan committed Oct 29, 2020
1 parent 28d9920 commit 33ea97c
Showing 1 changed file with 9 additions and 8 deletions.
17 changes: 9 additions & 8 deletions pytorch_lightning/trainer/connectors/checkpoint_connector.py
Original file line number Diff line number Diff line change
Expand Up @@ -93,17 +93,18 @@ def restore(self, checkpoint_path: str, on_gpu: bool) -> bool:
`True` if restored successfully else `False`
"""

# Try to load checkpoint from `checkpoint_path`. If failed, do not restore checkpoint.
try:
# if on_gpu:
# checkpoint = torch.load(checkpoint_path)
# else:
# load on CPU first
checkpoint = pl_load(checkpoint_path, map_location=lambda storage, loc: storage)
except Exception as _:
# Try to read the checkpoint file at `checkpoint_path`. If not exist, do not restore checkpoint.
fs = get_filesystem(checkpoint_path)
if not fs.exists(checkpoint_path):
log.info(f'failed to load model from checkpoint:{checkpoint_path}')
return False

# if on_gpu:
# checkpoint = torch.load(checkpoint_path)
# else:
# load on CPU first
checkpoint = pl_load(checkpoint_path, map_location=lambda storage, loc: storage)

# load model state
model = self.trainer.get_model()

Expand Down

0 comments on commit 33ea97c

Please sign in to comment.