`LightningModule.to_torchscript()` does not transfer check_inputs to correct device

[pfeatherstone]

Bug description

When running LightningModule.to_torchscript() you can pass both example_inputs and check_inputs. example_inputs gets transferred to the correct device. check_inputs does not.

What version are you seeing the problem on?

v2.0

How to reproduce the bug

No response

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

More info

No response

cc @Borda

[awaelchli]

@pfeatherstone Thanks for reporting. I'm labeling this as a community item. It would be nice if you could paste a simple example so that it's easier for a community contributor to get started.

[rajathpatel23]

@awaelchli , can I work on this issue once we have the directions to reproduce this bug

[pfeatherstone]

Yeah sorry I forgot to submit. I will do that soon

[rajathpatel23]

@pfeatherstone , wanted to check if you got a chance to find an example to reproduce the error

[pfeatherstone]

Oh yeah I forgot about this. I'll try do it first thing tomorrow if I remember

[pfeatherstone]

@awaelchli @rajathpatel23

import os
import torch
import lightning.pytorch as pl

class RandomDataset(torch.utils.data.Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("valid_loss", loss)

    def test_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("test_loss", loss)

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)

class CallExport(pl.callbacks.Callback):
    def on_validation_end(self, trainer, pl_module: BoringModel):
        if trainer.is_global_zero:
            print("Running TorchScript")
            input1 = torch.randn(10, 32)
            input2 = torch.randn(20, 32)
            pl_module.to_torchscript('/tmp/model.pt', method='trace', example_inputs=(input1,), check_inputs=(input2,))
        trainer.strategy.barrier()

def run():
    train_data  = torch.utils.data.DataLoader(RandomDataset(32, 64), batch_size=2)
    val_data    = torch.utils.data.DataLoader(RandomDataset(32, 64), batch_size=2)

    model = BoringModel()
    trainer = pl.Trainer(
        accelerator="gpu",
        devices=1,
        callbacks=[CallExport()],
        default_root_dir=os.getcwd(),
        limit_train_batches=1,
        limit_val_batches=1,
        limit_test_batches=1,
        num_sanity_val_steps=0,
        max_epochs=1,
        enable_model_summary=False,
    )
    trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)

if __name__ == "__main__":
    run()

[asingh9530]

@awaelchli looks like this should not be any major change following are the changes I think needed to be done in module.py.

    @torch.no_grad()
    def to_torchscript(
        self,
        file_path: Optional[Union[str, Path]] = None,
        method: Optional[str] = "script",
        example_inputs: Optional[Any] = None,
        **kwargs: Any,
    ) -> Union[ScriptModule, Dict[str, ScriptModule]]:
       
        mode = self.training

        if method == "script":
            with _jit_is_scripting():
                torchscript_module = torch.jit.script(self.eval(), **kwargs)
        elif method == "trace":
            # if no example inputs are provided, try to see if model has example_input_array set
            if example_inputs is None:
                if self.example_input_array is None:
                    raise ValueError(
                        "Choosing method=`trace` requires either `example_inputs`"
                        " or `model.example_input_array` to be defined."
                    )
                example_inputs = self.example_input_array
            
            # adding support for check_inputs
            if 'check_inputs' in kwargs and kwargs['check_inputs'] is not None:
                kwargs['check_inputs'] = self._on_before_batch_transfer(kwargs['check_inputs'])
                kwargs['check_inputs'] = self._apply_batch_transfer_handler(kwargs['check_inputs'])

            # automatically send example inputs to the right device and use trace
            example_inputs = self._on_before_batch_transfer(example_inputs)
            example_inputs = self._apply_batch_transfer_handler(example_inputs)
            with _jit_is_scripting():
                torchscript_module = torch.jit.trace(func=self.eval(), example_inputs=example_inputs, **kwargs)
        else:
            raise ValueError(f"The 'method' parameter only supports 'script' or 'trace', but value given was: {method}")

        self.train(mode)

        if file_path is not None:
            fs = get_filesystem(file_path)
            with fs.open(file_path, "wb") as f:
                torch.jit.save(torchscript_module, f)

        return torchscript_module

This followed by an added test cases to check for check_inputs device assertion ? Let me know I can raise this PR right away.