Add assert

torchmd · Feb 16, 2024 · 93fe195 · 93fe195
1 parent 0844242
commit 93fe195
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torchmdnet/priors/atomref.py b/torchmdnet/priors/atomref.py
@@ -34,6 +34,7 @@ def __init__(self, max_z=None, dataset=None, trainable=False, enable=True):
         if max_z is None and dataset is None:
             raise ValueError("Can't instantiate Atomref prior, all arguments are None.")
         if dataset is None:
+            assert max_z is not None, "max_z must be provided if dataset is None."
             atomref = torch.zeros(max_z, 1)
         else:
             atomref = dataset.get_atomref()