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