Check that Mg has the right element number

chaidiscovery · Dec 4, 2024 · 0ddfc90 · 0ddfc90
1 parent 3b43b2a
commit 0ddfc90
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/test_inference_dataset.py b/tests/test_inference_dataset.py
@@ -52,3 +52,4 @@ def test_ions_parsing(tokenizer: AllAtomResidueTokenizer):
     chain = chains[0]
     assert chain.structure_context.num_atoms == 1
     assert chain.structure_context.atom_ref_charge == 2
+    assert chain.structure_context.atom_ref_element.item() == 12