When moving tensors to a device using .to(device), the module is moved to the device, while tensors withing the Transformation remain on the cpu.
Instead relevant tensors should be torch.nn.Parameter objects, as these are designed to be associated with a module and correctly moved to a device alongside the module. This approach should ensure that all transformations and operations perform correctly post-device transfer.
Description
When moving tensors to a device using .to(device), the module is moved to the device, while tensors withing the Transformation remain on the cpu.
Instead relevant tensors should be
torch.nn.Parameterobjects, as these are designed to be associated with a module and correctly moved to a device alongside the module. This approach should ensure that all transformations and operations perform correctly post-device transfer.