I’m trying to migrate some code from Google Colab to AWS Sagemaker. I have a training script and at the end of the training I evaluate results using some functions, here comes the problem. I have troubles using numba. I have this function for cosine similarity:

```
@jit(nopython=True)
def fast_cosine(u, v):
m = u.shape[0]
udotv = 0
u_norm = 0
v_norm = 0
for i in range(m):
if (np.isnan(u[i])) or (np.isnan(v[i])):
continue
udotv += u[i] * v[i]
u_norm += u[i] * u[i]
v_norm += v[i] * v[i]
u_norm = np.sqrt(u_norm)
v_norm = np.sqrt(v_norm)
if (u_norm == 0) or (v_norm == 0):
ratio = 1.0
else:
ratio = udotv / (u_norm * v_norm)
return ratio
```

That function works flawlessy in Google Colab, taking couple of minutes to complete. When I use the same in Sagemaker, it takes hours. I also tried to add `cache=True`

and to set `NUMBA_CACHE_DIR`

to a path surely writable, in my case `/opt/ml/model`

. I also set `NUMBA_DEBUG_CACHE`

but the only output I get is this:

```
[cache] index saved to '/opt/ml/model/code_ec9d94cfa840991722b08139d5ab141658f37feb/train.fast_cosine-199.py38.nbi'
[cache] data saved to '/opt/ml/model/code_ec9d94cfa840991722b08139d5ab141658f37feb/train.fast_cosine-199.py38.1.nbc'
```

This function takes ages on AWS Sagemaker, while it takes only couple of minutes in Google Colab (where it takes much more without using Numba).