Typo

torchmd · Feb 13, 2024 · 907a614 · 907a614
1 parent fc3bb6b
commit 907a614
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchmdnet/priors/atomref.py b/torchmdnet/priors/atomref.py
@@ -47,7 +47,7 @@ def __init__(self, max_z=None, dataset=None, trainable=False, enable=True):
         if atomref.ndim == 1:
             atomref = atomref.view(-1, 1)
         self.register_buffer("initial_atomref", atomref)
-        self.atomref = nn.Embedding(len(atomref), 1, freeze=trainable and not enable)
+        self.atomref = nn.Embedding(len(atomref), 1, _freeze=trainable and not enable)
         self.atomref.weight.data.copy_(atomref)
         self.enable = enable