Update docs for map_location (Lightning-AI#920)

* update docs for map location

* update return description

pytorch_lightning/core/lightning.py

@@ -1023,10 +1023,15 @@ def load_from_metrics(cls, weights_path, tags_csv, map_location=None):
                     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 @@ def __init__(self, hparams):
 
         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
         -------