def test_torchscript_input_output_trace():
    """Test that traced LightningModule forward works with example_inputs."""
    model = BoringModel()
    example_inputs = torch.randn(1, 32)
    script = model.to_torchscript(example_inputs=example_inputs, method="trace")
    assert isinstance(script, torch.jit.ScriptModule)

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

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