def test_torchscript_device(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)

    script = model.to_torchscript()
    assert next(script.parameters()).device == device
    script_output = script(model.example_input_array.to(device))
    assert script_output.device == device
Test that scripted module is on the correct device.
tests/tests_pytorch/models/test_torchscript.py