PyTorch Lightning move tensor to correct device in validation_epoch_end
I would like to create a new tensor in a validation_epoch_end method of a LightningModule. From the official docs (page 48) it is stated that we should avoid direct .cuda() or .to(device) calls:

There are no .cuda() or .to() calls. . . Lightning does these for you.

and we are encouraged to use type_as method to transfer to the correct device.
new_x = new_x.type_as(x.type())
However, in a step validation_epoch_end I do not have any tensor to copy device from(by type_as method) in a clean way.
My question is what should I do if I want to create a new tensor in this method and transfer it to the device where is the model?
The only thing I can think of is to find a tensor in the outputs dictionary but it feels kinda messy:
avg_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
output = self(self.__test_input.type_as(avg_loss))

Is there any clean way to achieve that?

did you check part 3.4 (page 34) in the doc you linked ?

LightningModules know what device they are on!  construct tensors on the device directly to avoid CPU-&gt;Device transfer

t = tensor.rand(2, 2).cuda()# bad
(self is lightningModule)t = tensor.rand(2,2, device=self.device)# good 

I had a similar issue to create tensors this helped me. I hope it will help you too.
