1.0.0rc4 Save to TorchScript self.eval() device error
🐛 Bug
Likely self.eval() sends the model to the CPU, resulting in saving to a TorchScript file fails.
To Reproduce
Run the code sample below.
Code sample

Script can also be downloaded from: https://gist.github.com/NumesSanguis/388b4cfab2a8945afa85e8b79cd0c794
Most relevant code (extended to_torchscript to support trace):
# ...

    # use TorchScript trace
    def to_torchscript(...):
        # ...
                print(f"\n\nExample inputs device: {example_inputs.device}\n")
                scripted_module = torch.jit.trace(func=self.eval(), example_inputs=example_inputs, **kwargs)
        # ...
        return scripted_module
#...

if __name__ == '__main__':
    # ...
    # UNTIL HERE WORKS FINE

    # save model as TorchScript using eval()
    model.to_torchscript(file_path="example.pt", torchscript_approach='trace')
CLICK ME for full code

from __future__ import annotations  # don't crash on non-imported libraries when type checking
from typing import Dict, Optional, Union
import torch
from torch import nn as nn
import pytorch_lightning as pl
from pytorch_lightning import Trainer
from torch.utils.data import Dataset, DataLoader
from torch.nn import functional as F


class MyModel(pl.LightningModule):
    def __init__(self):  # **kwargs  # sample_rate
        super().__init__()
        # self.save_hyperparameters()
        # 1 sec of audio
        self.input_layer = nn.Linear(8000, 400, bias=True)
        self.hidden_layer = nn.Linear(400, 128, bias=True)
        self.output_layer = nn.Linear(128, 3, bias=True)

        self.criterion = nn.CrossEntropyLoss()

    def forward(self, input):
        x = F.relu_(self.input_layer(input))
        x = F.relu_(self.hidden_layer(x))
        output = self.output_layer(x)  # torch.sigmoid()
        return output

    def calculate_loss(self, prediction, target):
        loss = self.criterion(prediction, target)
        return loss

    def training_step(self, batch, batch_idx):
        input, target = batch
        # !!! use first batch to create an example input !!!
        if self.example_input_array is None:
            # we only need 1 sample, not a whole batch
            self.example_input_array = input
        prediction = self(input)
        loss = self.calculate_loss(prediction, target)
        return loss

    def validation_step(self, batch, batch_idx):
        input, target = batch
        prediction = self(input)
        loss = self.calculate_loss(prediction, target)
        return loss

    def test_step(self, batch, batch_idx):
        input, target = batch
        prediction = self(input)
        loss = self.calculate_loss(prediction, target)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

    # TODO create pull request for this
    def to_torchscript(
            self, file_path: Optional[str] = None, torchscript_approach: Optional[str] = 'script',
            example_inputs: torch.Tensor = None, **kwargs
    ) -> Union[ScriptModule, Dict[str, ScriptModule]]:
        # training or eval/test?
        mode = self.training
        with torch.no_grad():
            if torchscript_approach == 'script':
                scripted_module = torch.jit.script(self.eval(), **kwargs)
            elif torchscript_approach == 'trace':
                if example_inputs is None:
                    example_inputs = self.example_input_array
                print(f"\n\nExample inputs device: {example_inputs.device}\n")
                scripted_module = torch.jit.trace(func=self.eval(), example_inputs=example_inputs, **kwargs)
            else:
                raise ValueError(f"torchscript_approach only supports 'script' or 'trace', but value given was:"
                                 f"{torchscript_approach}")
        # set back whether we were training or not
        self.train(mode)

        if file_path is not None:
            torch.jit.save(scripted_module, file_path)

        return scripted_module


# DATA
class SimpleDataset(Dataset):
    def __init__(self, sample_rate=8000):
        self.sample_rate = sample_rate

    def __len__(self):
        return 16

    def __getitem__(self, idx):
        # 0, 1 or 2
        target = torch.randint(0, 3, size=(1, )).squeeze()
        # size 8000/16000 of 0.0, 0.5, or 1.0
        input = torch.full((self.sample_rate,), (target.float()/2).item())
        # torch.empty(self.sample_rate,).fill_(target.float()/2)
        return input, target


class SimpleDatamodule(pl.LightningDataModule):
    def setup(self, stage: str = None):
        pass

    def train_dataloader(self):
        return DataLoader(SimpleDataset(), batch_size=4)

    def val_dataloader(self):
        return DataLoader(SimpleDataset(), batch_size=4)

    def test_dataloader(self):
        return DataLoader(SimpleDataset(), batch_size=4)


if __name__ == '__main__':
    sr = 8000
    checkpoint_location = "example.ckpt"
    # network
    model = MyModel()
    # data
    dm = SimpleDatamodule()
    # train
    trainer = Trainer(max_epochs=2, deterministic=True, gpus=1)  # gpus=1,
    trainer.fit(model, dm)
    # save
    trainer.save_checkpoint(checkpoint_location)
    # UNTIL HERE WORKS FINE

    # save model as TorchScript using eval()
    model.to_torchscript(file_path="example.pt", torchscript_approach='trace')


Expected behavior
Model and Tensor input both being evaluated at device GPU.
Actual behavior
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
CLICK ME for full Traceback

yes, even hidden code blocks!
$ python save_as_torchscript_device_error.py 
GPU available: True, used: True
TPU available: False, using: 0 TPU cores
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

  | Name         | Type             | Params
--------------------------------------------------
0 | input_layer  | Linear           | 3 M   
1 | hidden_layer | Linear           | 51 K  
2 | output_layer | Linear           | 387   
3 | criterion    | CrossEntropyLoss | 0     
/home/*user*/anaconda3/envs/pytorchlit10/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:45: UserWarning: The dataloader, val dataloader 0, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 16 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  warnings.warn(*args, **kwargs)
/home/*user*/anaconda3/envs/pytorchlit10/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:45: UserWarning: The dataloader, train dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 16 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  warnings.warn(*args, **kwargs)
Epoch 1: 100%|█████████████████████████████████████████████████████████████████████| 8/8 [00:00<00:00, 148.42it/s, loss=0.954, v_num=55]
                              

Example inputs device: cuda:0

Traceback (most recent call last):
  File "/home/*user*/anaconda3/envs/pytorchlit10/lib/python3.8/site-packages/torch/jit/__init__.py", line 692, in run_mod_and_filter_tensor_outputs
    outs = wrap_retval(mod(*_clone_inputs(inputs)))
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "<string>", line 3, in forward

      def addmm(self: Tensor, mat1: Tensor, mat2: Tensor, beta: number = 1.0, alpha: number = 1.0):
          return self + mat1.mm(mat2)
                 ~~~~~~~~~~~~~~~~~~~ <--- HERE

      def batch_norm(input : Tensor, running_mean : Optional[Tensor], running_var : Optional[Tensor], training : bool, momentum : float, eps : float) -> Tensor:
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!


During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "eval_gpu_test.py", line 130, in <module>
    model.to_torchscript(file_path="example.pt", torchscript_approach='trace')
  File "eval_gpu_test.py", line 71, in to_torchscript
    scripted_module = torch.jit.trace(func=self.eval(), example_inputs=example_inputs, **kwargs)
  File "/home/*user*/anaconda3/envs/pytorchlit10/lib/python3.8/site-packages/torch/jit/__init__.py", line 953, in trace
    return trace_module(func, {'forward': example_inputs}, None,
  File "/home/*user*/anaconda3/envs/pytorchlit10/lib/python3.8/site-packages/torch/jit/__init__.py", line 1118, in trace_module
    _check_trace([inputs], func, check_trace_method,
  File "/home/*user*/anaconda3/envs/pytorchlit10/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 15, in decorate_context
    return func(*args, **kwargs)
  File "/home/*user*/anaconda3/envs/pytorchlit10/lib/python3.8/site-packages/torch/jit/__init__.py", line 734, in _check_trace
    traced_outs = run_mod_and_filter_tensor_outputs(traced_func, inputs, 'trace')
  File "/home/*user*/anaconda3/envs/pytorchlit10/lib/python3.8/site-packages/torch/jit/__init__.py", line 696, in run_mod_and_filter_tensor_outputs
    raise TracingCheckError(*graph_diagnostic_info(),
torch.jit.TracingCheckError: Tracing failed sanity checks!
Encountered an exception while running the trace with test inputs.
Exception:
	The following operation failed in the TorchScript interpreter.
	Traceback of TorchScript (most recent call last):
	  File "<string>", line 3, in forward
	
	      def addmm(self: Tensor, mat1: Tensor, mat2: Tensor, beta: number = 1.0, alpha: number = 1.0):
	          return self + mat1.mm(mat2)
	                 ~~~~~~~~~~~~~~~~~~~ <--- HERE
	
	      def batch_norm(input : Tensor, running_mean : Optional[Tensor], running_var : Optional[Tensor], training : bool, momentum : float, eps : float) -> Tensor:
	RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!


Environment

CUDA:

GPU:

GeForce GTX 1080 Ti


available:         True
version:           10.2


Packages:

numpy:             1.19.1
pyTorch_debug:     False
pyTorch_version:   1.6.0
pytorch-lightning: 1.0.0rc4
tqdm:              4.50.1


System:

OS:                Linux (Ubuntu 18.04)
architecture:

64bit
ELF


processor:         x86_64
python:            3.8.5
version:           #119-Ubuntu SMP Tue Sep 8 12:30:01 UTC 2020



Additional context
Might it be related to one of these eval() changes in v0.10?:
https://pytorch-lightning.readthedocs.io/en/stable/CHANGELOG.html?highlight=eval#changed