diff --git a/atomic_physics/common.py b/atomic_physics/common.py
index 320a46a..792c013 100644
--- a/atomic_physics/common.py
+++ b/atomic_physics/common.py
@@ -209,9 +209,9 @@ def index(
             MJvec = self.MJ[lev]
             inds = np.logical_and(inds, MJvec == MJ)
 
-        inds = np.argwhere(inds)
+        inds = np.nonzero(inds)[0]
         if len(inds) == 1:
-            inds = int(inds)
+            inds = inds[0]
         return inds + self.levels[level]._start_ind
 
     def level(self, state: int):