Fixed tensor not put onto device

[rmchurch]

New to Lightning. Saw some examples of this error but didn't seem to be my case. Below is the relevant part of the module. I run this with trainer = pl.Trainer(gpus=1). However, I get the error:

RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument mat2 in method wrapper_bmm)

Looking closer, self.A is not on the GPU. I thought it would be placed there automatically since it is in init. How should I place it on there? Its a fixed matrix (not learnable parameters, but does need to be backpropagated through). I see I can use self.device to put tensors on the GPU, but I want this matrix to be placed on the GPU in the init, without placing it every training step.

class MyModule(pl.LightningModule):
    def __init__(self,xin,yin,dt=1e-7):
        super().__init__()
        self.NN = Unet(channels=2).float()
        A = create_fixed_matrix(xin,yin)
        self.A = torch.from_numpy(A)
        self.time = torch.tensor([0.,dt])

    def forward(self,t,f):
        tmp = self.NN(f)
        fnextML = torch.einsum('xygh,bgh->bxy', self.A, tmp)
        return fnextML
    
    def training_step(self, batch, batch_idx):
        f,fnext = batch
        fnextML = self(self.time,f)
        loss = F.mse_loss(fnext,fnextML)
        return loss

[rmchurch]

I see it now, for non-parameters, I have to use self.register_buffer

https://pytorch-lightning.readthedocs.io/en/latest/accelerators/accelerator_prepare.html