🐛[BUG]: Cannot load model with additional buffers registered

[simonbyrne]

Version

0.6.0

On which installation method(s) does this occur?

Docker

Describe the issue

I wanted to include the normalization coefficients in the saved model, so I attached them via

net.register_buffer("norm_coeffs", norm_coeffs)

However when I reload the checkpoint via

net = Module.from_checkpoint(filename)

I get the following error

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.10/dist-packages/modulus/models/module.py", line 363, in from_checkpoint
    model.load_state_dict(model_dict)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 2153, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for EDMPrecond:
        Unexpected key(s) in state_dict: "norm_coeffs". 

Minimum reproducible example

No response

Relevant log output

No response

Environment details

No response

[mnabian]

@loliverhennigh could you please take a look at this issue?