Updating a register_buffer in PyTorch?
Correct way to update a register_buffer in PyTorch
I'm trying to determine the recommended way to update a register buffer which preserves the buffer's attributes. One &quot;hacky&quot; way to do this is shown below:
import torch
import torch.nn as nn


class SomeModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer(&quot;A&quot;, torch.tensor(2.0))

    def forward(self, x):
        return x * self.A

    def update_buffer(self, v):
        A = getattr(self, &quot;A&quot;)
        setattr(self, &quot;A&quot;, torch.tensor(v, dtype=A.dtype, device=A.device))


if __name__ == &quot;__main__&quot;:
    model = SomeModule()
    model.to(torch.float64)
    print(model(1.0).dtype)
    model.update_buffer(3.0)
    print(model(1.0).dtype)

When updating many buffers this gets a bit messy. Is there a recommended / better method to accomplish this?

Generally:
self.A = new_tensor

To copy the dtype and device from the previous tensor:
self.A = new_tensor.type_as(self.A)

If the new tensor has the same shape as the old, or you want to overwrite the current tensor's data with a new value, you can do:
self.A[:] = new_value


Examples:
&gt;&gt;&gt; mod.A = torch.tensor([3, 2, 1], dtype=int, device=&quot;cuda&quot;)
&gt;&gt;&gt; mod.A
tensor([3, 2, 1], device='cuda:0')
&gt;&gt;&gt; mod.A = torch.Tensor([1.5, 2, 3.5]).type_as(mod.A)
&gt;&gt;&gt; mod.A
tensor([1, 2, 3], device='cuda:0')
&gt;&gt;&gt; mod.A[:] = torch.Tensor([3, 4, 5])
&gt;&gt;&gt; mod.A
tensor([3, 4, 5], device='cuda:0')
&gt;&gt;&gt; mod.A[:] = 7                      
&gt;&gt;&gt; mod.A
tensor([7, 7, 7], device='cuda:0')
&gt;&gt;&gt; mod.state_dict()
OrderedDict([('A', tensor([7, 7, 7], device='cuda:0'))])

