Repository: Lightning-AI/pytorch-lightning
File: pytorch_lightning/core/lightning.py (Python)

Code before:
            >>> 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:
            This LightningModule as a torchscript, regardless of whether file_path is
            defined or not.
        """

        mode = self.training
        with torch.no_grad():
            if method == 'script':
                torchscript_module = torch.jit.script(self.eval(), **kwargs)
            elif method == 'trace':
                # 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
                example_input_array_device = move_data_to_device(example_inputs, device=self.device)
                torchscript_module = torch.jit.trace(func=self.eval(), example_inputs=example_input_array_device,
                                                     **kwargs)
            else:
                raise ValueError(f"The 'method' parameter only supports 'script' or 'trace', but value given was:"
                                 f"{method}")
        self.train(mode)

        if file_path is not None:
            torch.jit.save(torchscript_module, file_path)

        return torchscript_module

    @property
    def hparams(self) -> Union[AttributeDict, dict, Namespace]:
        if not hasattr(self, "_hparams"):
            self._hparams = AttributeDict()
        return self._hparams

    @property
    def hparams_initial(self) -> AttributeDict:
        if not hasattr(self, "_hparams_initial"):
            return AttributeDict()
        # prevent any change
        return copy.deepcopy(self._hparams_initial)

    @hparams.setter
    def hparams(self, hp: Union[dict, Namespace, Any]):

Review comment (suggestion):
```suggestion
                example_inputs = move_data_to_device(example_inputs, device=self.device)
                torchscript_module = torch.jit.trace(func=self.eval(), example_inputs=example_inputs, **kwargs)
```

Code after:
            ...                                     example_inputs=torch.randn(1, 64)))  # doctest: +SKIP
            >>> os.path.isfile("model_trace.pt")  # doctest: +SKIP
            True

        Return:
            This LightningModule as a torchscript, regardless of whether file_path is
            defined or not.
        """

        mode = self.training
        with torch.no_grad():
            if method == 'script':
                torchscript_module = torch.jit.script(self.eval(), **kwargs)
            elif method == 'trace':
                # 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
                # automatically send example inputs to the right device and use trace
                example_inputs = self.transfer_batch_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}")
        self.train(mode)

        if file_path is not None:
            torch.jit.save(torchscript_module, file_path)

        return torchscript_module

    @property
    def hparams(self) -> Union[AttributeDict, dict, Namespace]:
        if not hasattr(self, "_hparams"):
            self._hparams = AttributeDict()
        return self._hparams

    @property
    def hparams_initial(self) -> AttributeDict:
        if not hasattr(self, "_hparams_initial"):
            return AttributeDict()
        # prevent any change
        return copy.deepcopy(self._hparams_initial)

    @hparams.setter
    def hparams(self, hp: Union[dict, Namespace, Any]):
        hparams_assignment_name = self.__get_hparams_assignment_variable()
        self._hparams_name = hparams_assignment_name
        self._set_hparams(hp)
        # this resolves case when user does not uses `save_hyperparameters` and do hard assignement in init
        if not hasattr(self, "_hparams_initial"):
            self._hparams_initial = copy.deepcopy(self._hparams)