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, rtol=1e-5, atol=1e-8), (
        f"Scripted output {script_output} does not match model output {model_output}."
    )
Test that scripted LightningModule forward works.
tests/tests_pytorch/models/test_torchscript.py