def test_torchscript_device_with_check_inputs(device_str):
    """Test that scripted module is on the correct device."""
    device = torch.device(device_str)
    model = BoringModel().to(device)
    model.example_input_array = torch.randn(5, 32)

    check_inputs = torch.rand(5, 32)

    script = model.to_torchscript(method="trace", check_inputs=check_inputs)
    assert isinstance(script, torch.jit.ScriptModule)
Test that scripted module is on the correct device.
tests/tests_pytorch/models/test_torchscript.py