Secondary Pytorch LightningModule-declared Tensors not migrating to correct device when called from other LightningModule.
🐛 Bug
Background

After working on a Seq2Seq model with attention using only LightingModules, I was still getting on_device errors.
This also occurred with PackedSequence objects but I think that is due to a PyTorch cpu-specific implementation, not Lightning.
This also occurred with functions returning tensors called from LightingModules, but less of an issue than LightningModules themselves exhibiting the error.
ExampeL
The following class for Attention required the .cuda() flag ( .to(DEVICE) ):

class Attention(pl.LightningModule):
    '''
    Attention is calculated using key, value and query from Encoder and decoder.
    Below are the set of operations you need to perform for computing attention:
        energy = bmm(key, query)
        attention = softmax(energy)
        context = bmm(attention, value)
    '''
    def __init__(self):
        super(Attention, self).__init__()

    def forward(self, query, key, value, lens):
        '''
        :param query :(N, context_size) Query is the output of LSTMCell from Decoder
        :param key: (N, T_max, key_size) Key Projection from Encoder per time step
        :param value: (N, T_max, value_size) Value Projection from Encoder per time step
        :param lens: (N, T) Length of key and value, used for binary masking
        :return output: Attended Context
        :return attention: Attention mask that can be plotted
        '''

        energy = torch.bmm(key, query.unsqueeze(2)).squeeze(2) # (N, T_max, key_size) * (N, context_size, 1) = (N, T_max, 1) -> (N, T_max)

        # binary masking for padded positions
        mask = torch.arange(key.size(1)).unsqueeze(0) >= lens.unsqueeze(1) # (1, T) >= (B, 1) -> (N, T_max)
        mask = mask.to(DEVICE)
        energy.masked_fill_(mask, -1e9) # (N, T_max)
        attention = nn.functional.softmax(energy, dim=1) # (N, T_max)
        output = torch.bmm(attention.unsqueeze(1), value).squeeze(1) # (N, T_max)

        return output, attention




The class is called in another LightningModule during sequence decoding, which is in turn called by the training /testing loop.


Without the .to(DEVICE) flag, we get:


  File "/content/11785-hw4-lightning/models.py", line 46, in forward
    energy.masked_fill_(mask, -1e9) # (N, T_max)
RuntimeError: Expected object of device type cuda but got device type cpu for argument #2 'mask' in call to _th_masked_fill_bool_


I will reproduce the error using the minimal example provided:
Please reproduce using the BoringModel
COLAB
To Reproduce
Use following BoringModel and post here:
Same link
To reproduce, try having torch return a CPU tensor with a separate LightningModule loaded as a submodule. The CPU tensor will not migrate to the correct device causing the above problem. In this case, I usetorch.arange()on an unsqueezed scalar (the loss of the Boring Model. This is a common use for masked Seq2Seq loss, but it is common to see CPU tensors returned in LightningModules in other use cases. I've noted how this was apparent during an Attention-based implementation in the notebook as well as here.
Expected behavior
Torch-returned CPU tensor in a secondary LightningModule should be migrated to the correct device (cuda) but does not, such that a .cuda() call is required to fix the issue.
Environment
Note: Bugs with code are solved faster ! Colab Notebook should be made public !


IDE: Please, use our python bug_report_model.py template.


Colab Notebook: Please copy and paste the output from our environment collection script (or fill out the checklist below manually).


You can get the script and run it with:
wget https://raw.githubusercontent.com/PyTorchLightning/pytorch-lightning/master/tests/collect_env_details.py
# For security purposes, please check the contents of collect_env_details.py before running it.
python collect_env_details.py


CUDA:

GPU:

Tesla V100-SXM2-16GB


available:         True
version:           10.1


Packages:

numpy:             1.19.4
pyTorch_debug:     True
pyTorch_version:   1.7.0+cu101
pytorch-lightning: 1.1.2
tqdm:              4.41.1


System:

OS:                Linux
architecture:

64bit



processor:         x86_64
python:            3.6.9
version:           #1 SMP Thu Jul 23 08:00:38 PDT 2020



Additional context
I feel this is very important to address as projects scale, since many projects rely both on on more than a single LightningModule, and rely on torch-returned (as opposed to declared (?)) tensors. Thanks!