fix cpu time

fzimmermann89 · May 30, 2021 · 1d423d7 · 1d423d7
1 parent 446ddb4
commit 1d423d7
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/idi/simulation/time.py b/idi/simulation/time.py
@@ -16,15 +16,15 @@ def _decaysum(a, t, tau):
     return x
 
 
-@numba.njit(parallel=True)
+@_numba.njit(parallel=True)
 def _psort(a, b):
     '''
     idx = np.argsort(a, axis=-1)
     a = a[np.arange(a.shape[0])[:, None], idx]
     b = b[np.arange(a.shape[0])[:, None], idx]
     '''
-    for i in numba.prange(len(a)):
-        ids = np.argsort(a[i])
+    for i in _numba.prange(len(a)):
+        ids = _np.argsort(a[i])
         a[i, :] = a[i][ids]
         b[i, :] = b[i][ids]