Hardware agnostic training

[yhl48]

Is there a way to transfer a newly defined tensor in a torch.nn.Module class method to the device recognised by Lightning other than doing x.to(y) or calling a register_buffer in the initialisation as explained in the link below?

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

This is an issue in cases where a tensor has to be defined in a method without access to the other tensors' device.

# before lightning
def forward(self, x):
    z = torch.Tensor(2, 3)
    z = z.cuda(0)

# with lightning
def forward(self, x):
    z = torch.Tensor(2, 3)
    z = z.to(x)

# what if there is no input argument x
def another_method(self):
    z = torch.Tensor(2, 3)