I’ve encountered an issue while trying to register buffers from another module to a custom made module while doing the same thing with parameters (registering parameters) the method behaved as expected.
Here is a smaple of my code
import torch import torchvision.models class MyModule(torch.nn.Module): def __init__(self): super(MyModule, self).__init__() resnet18_model = torchvision.models.resnet18() my_model = MyModule() print("Checking Buffers") for buffer_name, buffer in list(resnet18_model.named_buffers()): my_model.register_buffer(buffer_name.replace('.', '_'), buffer) resnet18_model.cuda() for buffer in list(my_model.buffers()): print(buffer.is_cuda) print("Checking Parameters") for param_name, param in list(resnet18_model.named_parameters()): my_model.register_parameter(param_name.replace('.', '_'), param) resnet18_model.cuda() for param in list(my_model.parameters()): print(param.is_cuda)
My expected behaviour of this code was as follows :
after loading the resnet18 module into cuda, and after registering its parameters and buffers into my module I expected both of them to be in cuda, as you can see i’ve checked if is_coda flag is on and my results where as follows:
Checking Buffers False False . . . Checking Parameters True True . . .
it seems as if the
register_buffer uses some sort of “copying” the input buffers and not copying its address (reference) to the memory, while
register_parameters is acting as expected.
Is it a bug? or i’m missing something?
Is there a different way of registering the buffers pointers into other modules?