Azerus:

So, the simpler workaround I found is to avoid registering my float64_tensor, and move it to the current device if needed during the forward pass. Plus I added hooks to the state_dict in order to correctly set or get this float64_tensor when saving or loading my module.


I think you found a pretty good (and thorough) solution here, thank you for sharing!
You might double check if it works with distributed if that is something you need, and it seems inside the documented API. In particular, the obvious alternative – override Module.to –  would also work but bear more risk from messing with PyTorch internals.
Best regards
Thomas