def test_torchscript_save_load(tmp_path, modelclass):
    """Test that scripted LightningModule is correctly saved and can be loaded."""
    model = modelclass()
    output_file = str(tmp_path / "model.pt")
    script = model.to_torchscript(file_path=output_file)
    loaded_script = torch.jit.load(output_file)
    assert torch.allclose(next(script.parameters()), next(loaded_script.parameters()))
Test that scripted LightningModule is correctly saved and can be loaded.
tests/tests_pytorch/models/test_torchscript.py