PyTorch RuntimeError t == DeviceType::CUDAINTERNAL ASSERT FAILED
A PyTorch Lightning model works perfectly well on CPU using this Trainer configuration:
trainer = Trainer(
    gpus=0,
    max_epochs=10,
    gradient_clip_val=2,
    callbacks=[pl.callbacks.progress.TQDMProgressBar(refresh_rate=5)],
)

trainer.fit(model)

But running the exact same model on GPU (by changing gpus=-1 or gpus=1 in the above code) triggers the following error:
RuntimeError: t == DeviceType::CUDAINTERNAL ASSERT FAILED
at &quot;../c10/cuda/impl/CUDAGuardImpl.h&quot;:24, please report a bug to PyTorch. 

The model is as follows:
class TorchModel(LightningModule):
    def __init__(self):
        super(TorchModel, self).__init__()
        self.cat_layers = ModuleList([TorchCatEmbedding(cat) for cat in columns_to_embed])
        self.num_layers = ModuleList([LambdaLayer(lambda x: x[:, idx:idx+1]) for _, idx in numeric_columns])
        self.ffo = TorchFFO(len(self.num_layers) + sum([embed_dim(l) for l in self.cat_layers]), y.shape[1])
        self.softmax = torch.nn.Softmax(dim=1)

    def forward(self, inputs):
        cats = [c(inputs) for c in self.cat_layers]
        nums = [n(inputs) for n in self.num_layers]
        concat = torch.cat(cats + nums, dim=1)
        out = self.ffo(concat)
        out = self.softmax(out)
        return out

    def training_step(self, train_batch, batch_idx):
        x, y = train_batch
        y_hat = self.forward(x)
        loss = cce(torch.log(torch.maximum(torch.tensor(1e-8), y_hat)), y.argmax(dim=1))
        acc = tm.functional.accuracy(y_hat.argmax(dim=1), y.argmax(dim=1))
        self.log(&quot;loss&quot;, loss)
        self.log(&quot;acc&quot;, acc, prog_bar=True)
        self.log(&quot;lr&quot;, self.scheduler.get_last_lr()[-1], prog_bar=True)
        return loss

with TorchCatEmbedding and TorchFFO being two sub-models.
Is there any way to solve this issue?
PyTorch version:
&gt;&gt;&gt; torch.__version__
'1.10.1+cu113'

Cuda information:
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 460.32.03    Driver Version: 460.32.03    CUDA Version: 11.2     |
|-------------------------------+----------------------+----------------------+


This was due to a torch.tensor() declaration that wasn't transferred to GPU in the training step:
def training_step(self, train_batch, batch_idx):
    x, y = train_batch
    y_hat = self.forward(x)
    loss = cce(torch.log(torch.maximum(torch.tensor(1e-8), y_hat)), y.argmax(dim=1))
    acc = tm.functional.accuracy(y_hat.argmax(dim=1), y.argmax(dim=1))
    self.log(&quot;loss&quot;, loss)
    self.log(&quot;acc&quot;, acc, prog_bar=True)
    self.log(&quot;lr&quot;, self.scheduler.get_last_lr()[-1], prog_bar=True)
    return loss

Changing this:
loss = cce(
    torch.log(torch.maximum(torch.tensor(1e-8), y_hat)),
    y.argmax(dim=1)
)

to this by adding .type_as(y_hat):
loss = cce(
    torch.log(torch.maximum(torch.tensor(1e-8).type_as(y_hat), y_hat)), 
    y.argmax(dim=1)
)

solved the issue.
