move example inputs to correct device when tracing module

pytorch_lightning/core/lightning.py

@@ -20,7 +20,7 @@
 import tempfile
 from abc import ABC
 from argparse import Namespace
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union, Mapping
 
 import torch
 from pytorch_lightning import _logger as log
@@ -38,6 +38,7 @@
     collect_init_args,
     get_init_args,
 )
+from pytorch_lightning.utilities.apply_func import move_data_to_device
 from torch import ScriptModule, Tensor
 from torch.nn import Module
 from torch.optim.optimizer import Optimizer
@@ -1539,7 +1540,7 @@ def to_onnx(self, file_path: str, input_sample: Optional[Tensor] = None, **kwarg
 
     def to_torchscript(
         self, file_path: Optional[str] = None, method: Optional[str] = 'script',
-            example_inputs: Optional[torch.Tensor] = None, **kwargs
+            example_inputs: Optional[Union[torch.Tensor, Tuple[torch.Tensor]]] = None, **kwargs
     ) -> Union[ScriptModule, Dict[str, ScriptModule]]:
         """
         By default compiles the whole model to a :class:`~torch.jit.ScriptModule`.
@@ -1576,6 +1577,9 @@ def to_torchscript(
             >>> model = SimpleModel()
             >>> torch.jit.save(model.to_torchscript(), "model.pt")  # doctest: +SKIP
             >>> os.path.isfile("model.pt")  # doctest: +SKIP
+            >>> torch.jit.save(model.to_torchscript(file_path="model_trace.pt", method='trace', # doctest: +SKIP
+            ...                                     example_inputs=torch.randn(1, 64)))  # doctest: +SKIP
+            >>> os.path.isfile("model_trace.pt")  # doctest: +SKIP
             True
 
         Return:
@@ -1591,9 +1595,13 @@ def to_torchscript(
                 # if no example inputs are provided, try to see if model has example_input_array set
                 if example_inputs is None:
                     example_inputs = self.example_input_array
+                # dicts are not supported, so show the user an error; not raising an error to show the original error
+                if isinstance(example_inputs, Mapping):
+                    log.error("`example_inputs` should be a Tensor or a tuple of Tensors,"
+                              f" but got {type(example_inputs)}.")
                 # automatically send example inputs to the right device and use trace
-                torchscript_module = torch.jit.trace(func=self.eval(), example_inputs=example_inputs.to(self.device),
-                                                     **kwargs)
+                example_inputs = move_data_to_device(example_inputs, device=self.device)
+                torchscript_module = torch.jit.trace(func=self.eval(), example_inputs=example_inputs, **kwargs)
             else:
                 raise ValueError(f"The 'method' parameter only supports 'script' or 'trace', but value given was:"
                                  f"{method}")