Bug: AttributeError: module 'torch.nn' has no attribute 'Buffer' in multiple files.

[Tar-ive]

When running the pretrain.py script with a recent version of PyTorch, the training fails with an AttributeError because the code incorrectly uses nn.Buffer(...) instead of the correct self.register_buffer(...) method. This bug occurs in the following files:

models/sparse_embedding.py

models/layers.py

models/hrm/hrm_act_v1.py

I have a fix for this and will be opening a pull request shortly.

[imoneoi]

What's your PyTorch version? nn.Buffer is supposed to be the new API

[Tar-ive]

Hi @imoneoi, thanks for looking into this.

To answer your question, I was running this on PyTorch 2.1.0.

You're right that PyTorch APIs evolve, but in this case, self.register_buffer() is the correct and established method for adding a buffer to an nn.Module. The nn.Buffer(...) constructor that was in the code doesn't exist, which was causing the AttributeError during initialization.

Switching all instances to self.register_buffer() fixed the bug and makes the code compatible with PyTorch 2.1 and newer versions. Hope this helps clarify!