# Buggy code
def to_onnx(self, file_path: str, input_sample: Optional[Tensor] = None, **kwargs):
    """Saves the model in ONNX format

    Args:
        file_path: The path of the file the model should be saved to.
        input_sample: A sample of an input tensor for tracing.
        **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))
        ...     model.to_onnx(tmpfile.name, input_sample, export_params=True)
        ...     os.path.isfile(tmpfile.name)
        True
    """

    if isinstance(input_sample, Tensor):
        input_data = input_sample
    elif self.example_input_array is not None:
        input_data = self.example_input_array
    else:
        raise ValueError(
            "`input_sample` and `example_input_array` tensors are both missing."
        )

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

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


# Traceback (RuntimeError)
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-32-cd8009a0b6a3> in <module>
1 filepath = 'model.onnx'
----> 2 model.to_onnx(filepath, export_params=True)

/opt/conda/lib/python3.7/site-packages/pytorch_lightning/core/lightning.py in to_onnx(self, file_path, input_sample, **kwargs)
1721         if 'example_outputs' not in kwargs:
1722             self.eval()
-> 1723             kwargs['example_outputs'] = self(input_data)
1724
1725         torch.onnx.export(self, input_data, file_path, **kwargs)

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
548             result = self._slow_forward(*input, **kwargs)
549         else:
--> 550             result = self.forward(*input, **kwargs)
551         for hook in self._forward_hooks.values():
552             hook_result = hook(self, input, result)

<ipython-input-24-51cae3b5e57f> in forward(self, inputs)
20
21     def forward(self, inputs):
---> 22         return self.model(inputs)
23
24     def training_step(self, batch, batch_idx):

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
548             result = self._slow_forward(*input, **kwargs)
549         else:
--> 550             result = self.forward(*input, **kwargs)
551         for hook in self._forward_hooks.values():
552             hook_result = hook(self, input, result)

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/container.py in forward(self, input)
98     def forward(self, input):
99         for module in self:
--> 100             input = module(input)
101         return input
102

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/module.py in __call__(self, *input, **kwargs)
548             result = self._slow_forward(*input, **kwargs)
549         else:
--> 550             result = self.forward(*input, **kwargs)
551         for hook in self._forward_hooks.values():
552             hook_result = hook(self, input, result)

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/conv.py in forward(self, input)
351
352     def forward(self, input):
--> 353         return self._conv_forward(input, self.weight)
354
355 class Conv3d(_ConvNd):

/opt/conda/lib/python3.7/site-packages/torch/nn/modules/conv.py in _conv_forward(self, input, weight)
348                             _pair(0), self.dilation, self.groups)
349         return F.conv2d(input, weight, self.bias, self.stride,
--> 350                         self.padding, self.dilation, self.groups)
351
352     def forward(self, input):

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same

# Fixed code
def to_onnx(self, file_path: str, input_sample: Optional[Tensor] = None, **kwargs):
    """Saves the model in ONNX format

    Args:
        file_path: The path of the file the model should be saved to.
        input_sample: A sample of an input tensor for tracing.
        **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))
        ...     model.to_onnx(tmpfile.name, input_sample, export_params=True)
        ...     os.path.isfile(tmpfile.name)
        True
    """

    if isinstance(input_sample, Tensor):
        input_data = input_sample
    elif self.example_input_array is not None:
        input_data = self.example_input_array
    else:
        if input_sample is not None:
            raise ValueError(
                f"Received `input_sample` of type {type(input_sample)}. Expected type is `Tensor`"
            )
        else:
            raise ValueError(
                "Could not export to ONNX since neither `input_sample` nor"
                " `model.example_input_array` attribute is set."
            )
    input_data = input_data.to(self.device)
    if "example_outputs" not in kwargs:
        self.eval()
        with torch.no_grad():
            kwargs["example_outputs"] = self(input_data)

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