torch.nn.Module register

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

torch.nn.Module register_buffer()

Recommendations