diff --git a/mp_nerf/kb_proteins.py b/mp_nerf/kb_proteins.py
index 5424f2e..0dc36b5 100644
--- a/mp_nerf/kb_proteins.py
+++ b/mp_nerf/kb_proteins.py
@@ -721,12 +721,11 @@ def make_torsion_mask(aa, fill=False):
     """ Gives the dihedral of the bond originating each atom. """
     mask = np.zeros(14)
     # backbone
-    if aa == "_":
-        return aa
+
     # sidechain
     for i, torsion in enumerate(SC_BUILD_INFO[aa]['torsion-vals']):
         if fill: 
-            mask[4+i] = MP2SC_INFO[aa][ SC_BUILD_INFO[aa]["atom-names"] ]["bond_dihedral"]
+            mask[4+i] = MP2SC_INFO[aa][ SC_BUILD_INFO[aa]["atom-names"][i] ]["bond_dihedral"]
         else: 
             # https://github.com/jonathanking/sidechainnet/blob/master/sidechainnet/structure/StructureBuilder.py#L372
             # 999 is an anotation -- change later 
diff --git a/tests/tes_main.py b/tests/tes_main.py
index 8b420f7..8618559 100644
--- a/tests/tes_main.py
+++ b/tests/tes_main.py
@@ -1,6 +1,11 @@
 import numpy as np
 import torch
 
+from mp_nerf import *
+from mp_nerf.utils import *
+from mp_nerf.kb_proteins import *
+from mp_nerf.proteins import *
+
 def test_nerf_and_dihedral():
     # create points
     a = torch.tensor([1,2,3]).float()
@@ -26,5 +31,5 @@ def test_nerf_and_dihedral():
     # reconstruct
     # doesnt work because the scn angle was not measured correctly
     # so the method corrects that incorrection
-    assert (nerf_torch(a, b, c, l, theta, chi - np.pi) - torch.tensor([1,0,6])).sum().abs() < 0.1
-    assert get_dihedral_torch(a, b, c, d).item() == chi
\ No newline at end of file
+    assert (mp_nerf_torch(a, b, c, l, theta, chi - np.pi) - torch.tensor([1,0,6])).sum().abs() < 0.1
+    assert get_dihedral(a, b, c, d).item() == chi
\ No newline at end of file