Thanks for the feedback.
Good point about the distributed. I’m not using it, but I like to think of most use case (unfortunatly, I have only one gpu so i don’t know if I can actually perform the test).
I also though about overwriting Module.to but it would not work if my custom_module is registered as sub-module of model when running model.to(dtype), as only the to() method from model would be called and the not the one from custom_module. To do such approach I would have to overwrite the .to method from the float64_tensor itself, if I understood the source correctly (see the line return self._apply(convert) and how the function convert is defined).