Repository: Lightning-AI/pytorch-lightning
File: tests/models/test_torchscript.py (Python)

Code before:
# limitations under the License.
from distutils.version import LooseVersion

import pytest
import torch

from tests.base import BoringModel
from tests.base.datamodules import TrialMNISTDataModule
from tests.base.models import ParityModuleRNN, BasicGAN


@pytest.mark.parametrize("modelclass", [
    BoringModel,
    ParityModuleRNN,
    BasicGAN,
])
def test_torchscript_input_output(modelclass):
    """ Test that scripted LightningModule forward works. """
    model = modelclass()
    script = model.to_torchscript()
    assert isinstance(script, torch.jit.ScriptModule)

    model.eval()
    with torch.no_grad():
        model_output = model(model.example_input_array)
        script_output = script(model.example_input_array)
    assert torch.allclose(script_output, model_output)


@pytest.mark.parametrize("modelclass", [
    BoringModel,
    ParityModuleRNN,
    BasicGAN,
])
def test_torchscript_example_input_output_trace(modelclass):
    """ Test that traced LightningModule forward works with example_input_array """
    model = modelclass()
    script = model.to_torchscript(method='trace')
    assert isinstance(script, torch.jit.ScriptModule)

    model.eval()
    with torch.no_grad():
        model_output = model(model.example_input_array)
        script_output = script(model.example_input_array)
    assert torch.allclose(script_output, model_output)


def test_torchscript_input_output_trace():
    """ Test that traced LightningModule forward works with example_inputs """
    model = BoringModel()
    example_inputs = torch.randn(1, 32)

Review comment (suggestion):
```suggestion
        model_output = model(model.example_input_array)
    script_output = script(model.example_input_array)
```
this should be sufficient

Code after:
# limitations under the License.
from distutils.version import LooseVersion

import pytest
import torch

from tests.base import BoringModel
from tests.base.datamodules import TrialMNISTDataModule
from tests.base.models import ParityModuleRNN, BasicGAN


@pytest.mark.parametrize("modelclass", [
    BoringModel,
    ParityModuleRNN,
    BasicGAN,
])
def test_torchscript_input_output(modelclass):
    """ Test that scripted LightningModule forward works. """
    model = modelclass()

    if isinstance(model, BoringModel):
        model.example_input_array = torch.randn(5, 32)

    script = model.to_torchscript()
    assert isinstance(script, torch.jit.ScriptModule)

    model.eval()
    with torch.no_grad():
        model_output = model(model.example_input_array)

    script_output = script(model.example_input_array)
    assert torch.allclose(script_output, model_output)


@pytest.mark.parametrize("modelclass", [
    BoringModel,
    ParityModuleRNN,
    BasicGAN,
])
def test_torchscript_example_input_output_trace(modelclass):
    """ Test that traced LightningModule forward works with example_input_array """
    model = modelclass()

    if isinstance(model, BoringModel):
        model.example_input_array = torch.randn(5, 32)

    script = model.to_torchscript(method='trace')
    assert isinstance(script, torch.jit.ScriptModule)

    model.eval()
    with torch.no_grad():