def test_torchscript_example_input_output_trace(modelclass):
    """Test that traced LightningModule forward works with example_input_array."""
    torch.manual_seed(1)
    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():
        model_output = model(model.example_input_array)

    script_output = script(model.example_input_array)
    assert torch.allclose(script_output, model_output)
Test that traced LightningModule forward works with example_input_array.
tests/tests_pytorch/models/test_torchscript.py