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

Code before:
            ...     def forward(self, x):
            ...         return torch.relu(self.l1(x.view(x.size(0), -1)))

            >>> with tempfile.NamedTemporaryFile(suffix='.onnx', delete=False) as tmpfile:
            ...     model = SimpleModel()
            ...     input_sample = torch.randn((1, 64))
            ...     model.to_onnx(tmpfile.name, input_sample, export_params=True)
            ...     os.path.isfile(tmpfile.name)
            True
        """
        if input_sample is None:
            if self.example_input_array is None:
                raise ValueError(
                    "Could not export to ONNX since neither `input_sample` nor"
                    " `model.example_input_array` attribute is set."
                )
            input_sample = self.example_input_array

        if isinstance(input_sample, Tensor):
            input_sample = (input_sample,)
        elif not (isinstance(input_sample, tuple) and all(isinstance(inp, Tensor) for inp in input_sample)):
            raise ValueError(
                "Could not export to ONNX since input_sample is neither a Tensor nor tuple of Tensors"
            )

        input_sample = self.transfer_batch_to_device(input_sample, self.device)

        if "example_outputs" not in kwargs:
            self.eval()
            kwargs["example_outputs"] = self(*input_sample)

        torch.onnx.export(self, input_sample, file_path, **kwargs)

    @torch.no_grad()
    def to_torchscript(
        self, file_path: Optional[str] = None, method: Optional[str] = 'script',
        example_inputs: Optional[Union[Tensor, Tuple[Tensor]]] = None, **kwargs
    ) -> Union[ScriptModule, Dict[str, ScriptModule]]:
        """
        By default compiles the whole model to a :class:`~torch.jit.ScriptModule`.
        If you want to use tracing, please provided the argument `method='trace'` and make sure that either the
        example_inputs argument is provided, or the model has self.example_input_array set.
        If you would like to customize the modules that are scripted you should override this method.
        In case you want to return multiple modules, we recommend using a dictionary.

        Args:
            file_path: Path where to save the torchscript. Default: None (no file saved).
            method: Whether to use TorchScript's script or trace method. Default: 'script'
            example_inputs: An input tensor or tuple of tensors to be used to do tracing when method is set to 'trace'.
              Default: None (Use self.example_input_array)
            **kwargs: Additional arguments that will be passed to the :func:`torch.jit.script` or

Review comment (performance):
quick question: now that it is restricts that input_sample must be a tensor or tuple of tensors, should we use `self.transfer_batch_to_device` or `move_data_to_device`? Since there is a possibility that one might have override `transfer_batch_to_device` to adapt some custom batch type like dict or something else.

Code after:
                hp = {arg: init_args[arg] for arg in args if isinstance(arg, str)}
                self._hparams_name = "kwargs"

        # `hparams` are expected here
        if hp:
            self._set_hparams(hp)
        # make deep copy so  there is not other runtime changes reflected
        self._hparams_initial = copy.deepcopy(self._hparams)

    def _set_hparams(self, hp: Union[dict, Namespace, str]) -> None:
        if isinstance(hp, Namespace):
            hp = vars(hp)
        if isinstance(hp, dict):
            hp = AttributeDict(hp)
        elif isinstance(hp, PRIMITIVE_TYPES):
            raise ValueError(f"Primitives {PRIMITIVE_TYPES} are not allowed.")
        elif not isinstance(hp, ALLOWED_CONFIG_TYPES):
            raise ValueError(f"Unsupported config type of {type(hp)}.")

        if isinstance(hp, dict) and isinstance(self.hparams, dict):
            self.hparams.update(hp)
        else:
            self._hparams = hp

    @torch.no_grad()
    def to_onnx(
        self,
        file_path: Union[str, Path],
        input_sample: Optional[Any] = None,
        **kwargs,
    ):
        """
        Saves the model in ONNX format

        Args:
            file_path: The path of the file the onnx model should be saved to.
            input_sample: An input for tracing. Default: None (Use self.example_input_array)
            **kwargs: Will be passed to torch.onnx.export function.

        Example:
            >>> class SimpleModel(LightningModule):
            ...     def __init__(self):
            ...         super().__init__()
            ...         self.l1 = torch.nn.Linear(in_features=64, out_features=4)
            ...
            ...     def forward(self, x):
            ...         return torch.relu(self.l1(x.view(x.size(0), -1)))

            >>> with tempfile.NamedTemporaryFile(suffix='.onnx', delete=False) as tmpfile:
            ...     model = SimpleModel()
            ...     input_sample = torch.randn((1, 64))