PyTorch Lightning move tensor to correct device in validation_epoch_end

Question

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?

Answer 1

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->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.