I feel like that is not what @ptrblck was saying here: Multi GPU backwards hook on wrong device. It sounded like doing register_buffer allows you to define “self” variables that ensure that they work properly on multiple GPUs. Did I misread his post about that?
I feel like I need to use register_backward_hook. It is also working entirely correct for me on one GPU. Is there a page or thread you can point me to for what to do instead?