def test_torchscript_save_load_custom_filesystem(tmp_path, modelclass):
    """Test that scripted LightningModule is correctly saved and can be loaded with custom filesystems."""
    _DUMMY_PRFEIX = "dummy"
    _PREFIX_SEPARATOR = "://"

    class DummyFileSystem(LocalFileSystem): ...

    fsspec.register_implementation(_DUMMY_PRFEIX, DummyFileSystem, clobber=True)

    model = modelclass()
    output_file = os.path.join(_DUMMY_PRFEIX, _PREFIX_SEPARATOR, tmp_path, "model.pt")
    script = model.to_torchscript(file_path=output_file)

    fs = get_filesystem(output_file)
    with fs.open(output_file, "rb") as f:
        loaded_script = torch.jit.load(f)

    assert torch.allclose(next(script.parameters()), next(loaded_script.parameters()))
Test that scripted LightningModule is correctly saved and can be loaded with custom filesystems.
tests/tests_pytorch/models/test_torchscript.py