torch.nn.Module register_buffer()

Creator
Creator
Seonglae Cho
Created
Created
2025 Apr 10 20:11
Editor
Edited
Edited
2025 Apr 10 20:14
Refs
Refs
Register
torch.Tensor
into
Torch state_dict
but to not appears in learnable parameters
torch Model.parameters()
. It is used for registering constant configuration tensor or initialization tensor for the model should be managed by a Tensor in the same device.
import torch class MyModel(torch.nn.Module): def __init__(self): super().__init__() self.register_buffer("mask", torch.ones(10)) def forward(self, x): return x * self.mask # Example usage model = MyModel() print(model.mask) # Persistent buffer, not a parameter
 
 
 
 
 
 
 

Recommendations