RuntimeError: All input tensors must be on the same device. Received cpu and cuda:0
🐛 Bug
At the end of the epoch, I get the error mentioned in the title. Here is a full stack-trace:
File "main.py", line 255, in <module>
    trainer.fit(trainVQG, data_loader, val_data_loader)
  File "/data/nv419/anaconda3/envs/blt-vqg/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 473, in fit
    results = self.accelerator_backend.train()
  File "/data/nv419/anaconda3/envs/blt-vqg/lib/python3.8/site-packages/pytorch_lightning/accelerators/cpu_accelerator.py", line 60, in train
    results = self.train_or_test()
  File "/data/nv419/anaconda3/envs/blt-vqg/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 69, in train_or_test
    results = self.trainer.train()
  File "/data/nv419/anaconda3/envs/blt-vqg/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 524, in train
    self.train_loop.run_training_epoch()
  File "/data/nv419/anaconda3/envs/blt-vqg/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 637, in run_training_epoch
    self.run_on_epoch_end_hook(epoch_output)
  File "/data/nv419/anaconda3/envs/blt-vqg/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 868, in run_on_epoch_end_hook
    self.trainer.logger_connector.on_train_epoch_end()
  File "/data/nv419/anaconda3/envs/blt-vqg/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/logger_connector/logger_connector.py", line 370, in on_train_epoch_end
    self.cached_results.has_batch_loop_finished = True
  File "/data/nv419/anaconda3/envs/blt-vqg/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/logger_connector/epoch_result_store.py", line 445, in has_batch_loop_finished
    self.auto_reduce_results_on_epoch_end()
  File "/data/nv419/anaconda3/envs/blt-vqg/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/logger_connector/epoch_result_store.py", line 435, in auto_reduce_results_on_epoch_end
    hook_result.auto_reduce_results_on_epoch_end()
  File "/data/nv419/anaconda3/envs/blt-vqg/lib/python3.8/site-packages/pytorch_lightning/trainer/connectors/logger_connector/epoch_result_store.py", line 229, in auto_reduce_results_on_epoch_end
    opt_outputs = time_reduced_outputs[0].__class__.reduce_on_epoch_end(time_reduced_outputs)
  File "/data/nv419/anaconda3/envs/blt-vqg/lib/python3.8/site-packages/pytorch_lightning/core/step_result.py", line 519, in reduce_on_epoch_end
    recursive_stack(result)
  File "/data/nv419/anaconda3/envs/blt-vqg/lib/python3.8/site-packages/pytorch_lightning/core/step_result.py", line 660, in recursive_stack
    result[k] = collate_tensors(v)
  File "/data/nv419/anaconda3/envs/blt-vqg/lib/python3.8/site-packages/pytorch_lightning/core/step_result.py", line 682, in collate_tensors
    return torch.stack(items)
RuntimeError: All input tensors must be on the same device. Received cpu and cuda:0

The error gets thrown at the end of an epoch. The model I've built is a VAE/Latent Variable model. Note that this error does NOT get thrown (and the code works perfectly fine) if I do not use the VAE part of the model. The BoringModel doesn't reproduce the error, and I've got some research code in here which I'd rather not make public at this point in time. However, these are the parts of my code which may be contributing to the error:
# training_step
def training_step(self, batch, batch_idx):
        if self.args.num_warmup_steps < self.iter:
            self.latent_transformer = True
            self.model.latent_transformer = True
            self.configure_optimizers()  # reset the momentum

        loss, kld = self(batch)
        total_loss, loss_rec, loss_kl = self.calculate_losses(loss, kld)
        self.log('total train loss', total_loss)
        self.log('rec train loss', loss_rec)
        self.log('kl train loss', loss_kl)
        self.iter += 1
        return total_loss
# calculating losses
    def calculate_losses(self, loss, kld, r=0.5):
        if kld is None:
            loss_rec = loss
            total_loss = loss
            loss_kl = torch.tensor(0)
        else:
            cycle_num = (self.args.total_training_steps/4)
            mod = self.iter % cycle_num
            temp = mod/cycle_num
            beta = 1
            if temp <= r:
                beta = 1/(1 + np.exp(-temp))

            loss_kl = kld
            loss_rec = loss
            total_loss = loss + beta * kld

        return total_loss.to(self.args.device), loss_rec, loss_kl
# inside self.latent
class LatentNorm(nn.Module):
    def __init__(self, args):
        super(LatentNorm, self).__init__()
        self.args = args

        self.latent_encoder = nn.Sequential(
            nn.Linear(args.hidden_dim, args.latent_dim*2),
            nn.LeakyReLU(0.2),
            nn.Linear(args.latent_dim*2, args.latent_dim*2)
        )

    def forward(self, hidden_states):
        latents = self.latent_encoder(hidden_states)
        mu, logvar = latents[:, :, :self.args.latent_dim], latents[:, :, self.args.latent_dim:]
        std = torch.exp(0.5*logvar)
        eps = torch.randn_like(std)
        z = eps.mul(std).add_(mu)
        kld = gaussian_kld_norm(mu, logvar)
        return z, kld
# inside model
kld = None
if self.latent_transformer:
    encoder_hidden_states, kld = self.latent(encoder_hidden_states)
Interestingly, OP of this #5053 issue also seemed to be using a VAE and had the error thrown (although the solution to his problem isn't the solution to my problem)
Expected behavior
For the model to continue training as normal
Environment

CUDA:
- GPU:
- Tesla V100-PCIE-32GB
- Tesla V100-PCIE-32GB
- GeForce RTX 2080 Ti
- GeForce RTX 2080 Ti
- GeForce RTX 2080 Ti
- GeForce RTX 2080 Ti
- GeForce RTX 2080 Ti
- GeForce RTX 2080 Ti
- GeForce RTX 2080 Ti
- GeForce RTX 2080 Ti
- available:         True
- version:           10.2
Packages:
- numpy:             1.19.2
- pyTorch_debug:     False
- pyTorch_version:   1.7.1
- pytorch-lightning: 1.1.4
- tqdm:              4.60.0
System:
- OS:                Linux
- architecture:
- 64bit
- ELF
- processor:         x86_64
- python:            3.8.8
- version:           #59~18.04.1-Ubuntu SMP Wed Oct 21 12:14:56 UTC 2020

I'm running this code on a V100