How to pass additional parameters to training_step/validation step?

[laughingrice]

I have a loss function that depends on an additional tensor.

Something similar to:

x, y = batch
m = self(x)

m = A @ m
loss = mse_loss(m, y)

I tried adding A to the class as a hyperparameter (not great), but that crashes the training on checkpointing saying that not all tensors are on the same device (found at least two devices, cuda:2 and cpu!)
It also doesn't let me run on multiple GPUs as A is a sparse matrix and won't pickle

Is there a way to pass additional parameters to the training_step/validation step? maybe override torch.utils.data.DataLoader to pass a triplet instead of a doublet? - how would I go about taking care of passing things to the right devices for multi GPU training?