How to use a loss function on GPU

[fishbotics]

Hi all,

I have a loss function which is a callable instance of a class. The class itself has some state which is stored on the GPU. Basically the class has some tensors that it applies to project the network's output into some space before computing the loss in the projected subspace. Since these tensors need to live on the same device as the LightningModule, is there any special way to ensure this happens? I got rid of all my calls to .to() as specified in the docs, but now I'm not sure how to make sure this loss class is transferred to the correct device when the model is as well.

Thanks!

[awaelchli]

You can make them buffers.
https://pytorch-lightning.readthedocs.io/en/latest/advanced/multi_gpu.html#init-tensors-using-type-as-and-register-buffer

[fishbotics]

This is a good idea! Unfortunately, in my case, it's more complicated that just a single projection tensor that I could store. This is for a robotics task and I'm using a differentiable forward kinematics model (i.e. a chain of tensors based on the kinematics of the specific robot).

What I ended up doing (which seems to work if I understand correctly) is that I initialize the loss function in a setup() method on my module and so the loss function class is initialized per gpu.