From c56ee8bdee0582a6ae10bcbd8dccbeac2e7eb09e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Adrian=20W=C3=A4lchli?= Date: Sun, 23 Feb 2020 21:01:08 +0100 Subject: [PATCH] Update docs for map_location (#920) * update docs for map location * update return description --- pytorch_lightning/core/lightning.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py index 8790946b..7a09aabb 100644 --- a/pytorch_lightning/core/lightning.py +++ b/pytorch_lightning/core/lightning.py @@ -1023,10 +1023,15 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): drop_prob,0.2 batch_size,32 - map_location (dict): A dictionary mapping saved weight GPU devices to new - GPU devices (example: {'cuda:1':'cuda:0'}) + map_location (dict | str | torch.device | function): + If your checkpoint saved a GPU model and you now load on CPUs + or a different number of GPUs, use this to map to the new setup + (example: {'cuda:1':'cuda:0'}). + The behaviour is the same as in + `torch.load `_. + Return: - LightningModule with loaded weights + LightningModule with loaded weights and hyperparameters (if available). Example ------- @@ -1097,11 +1102,14 @@ class LightningModule(ABC, GradInformation, ModelIO, ModelHooks): Args: checkpoint_path (str): Path to checkpoint. - map_location (dic): If your checkpoint saved from a GPU model and you now load on CPUs + map_location (dict | str | torch.device | function): + If your checkpoint saved a GPU model and you now load on CPUs or a different number of GPUs, use this to map to the new setup. + The behaviour is the same as in + `torch.load `_. Return: - LightningModule with loaded weights. + LightningModule with loaded weights and hyperparameters (if available). Example -------