Thanks for taking a look into the issue.
I’m not storing these tensors manually and I can’t detach them in the training step as PyTorch Lightning needs it for backward. I tried detaching, using loss.item() in comet_logs = {}. Nothing worked so far.
I’ll look around in PyTorch Lightning code for details what they are doing with these dicts.
But any ideas would be helpful.