wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-02 12:40:56 +08:00
Code Issues Packages Projects Releases Wiki Activity

Update docs for map_location (#920)

Browse Source 
* update docs for map location

* update return description
This commit is contained in:
Adrian Wälchli
2020-02-23 15:01:08 -05:00
committed by GitHub
parent 5778a4131c
commit c56ee8bdee
1 changed files with 13 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/core/lightning.py
+13 -5
View File
@@ -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 <https://pytorch.org/docs/stable/torch.html#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 <https://pytorch.org/docs/stable/torch.html#torch.load>`_.
Return:
LightningModule with loaded weights.
LightningModule with loaded weights and hyperparameters (if available).
Example
-------