`06_precompute_cache.py` fails when running in cpu mode

The put_along_axis call at line 70:

sense2vec/scripts/06_precompute_cache.py

Line 70 in 61eb3a7

xp.put_along_axis(sims, indices, -xp.inf, axis=1)

If I change from the numpy call xp.put_along_axis to the put_along_axis defined in the script, this will fail because I dont have cupy installed (because this is running on my macbook pro which does not have nvidia gpu and hence no cuda installed).

Is there any easy change I can make to the put_along_axis function to get cpu mode working? Unfortunately I don't understanding enough about what is going on in that function to make the required change.

josh@JoshsMacBook ~/sense2vec (master=)
$ git pull origin master
From https://github.com/explosion/sense2vec
 * branch            master     -> FETCH_HEAD
Already up to date.
josh@JoshsMacBook ~/sense2vec (master=)
$ python ./scripts/06_precompute_cache.py -c 1000 -n 10 $S2V_MODEL_PATH
✔ Loaded 1,195,261 vectors with dimension 128
✔ Normalized (mean 3.76, variance 1.89)
ℹ Finding 10 neighbors among 1,000 most frequent
  0%|                                                                                               | 0/1168 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "./scripts/06_precompute_cache.py", line 176, in <module>
    plac.call(main)
  File "/usr/local/lib/python3.7/site-packages/plac_core.py", line 328, in call
    cmd, result = parser.consume(arglist)
  File "/usr/local/lib/python3.7/site-packages/plac_core.py", line 207, in consume
    return cmd, self.func(*(args + varargs + extraopts), **kwargs)
  File "./scripts/06_precompute_cache.py", line 70, in main
    xp.put_along_axis(sims, indices, -xp.inf, axis=1)
  File "<__array_function__ internals>", line 6, in put_along_axis
  File "/usr/local/lib/python3.7/site-packages/numpy/lib/shape_base.py", line 262, in put_along_axis
    arr[_make_along_axis_idx(arr_shape, indices, axis)] = values
IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (1024,1) (1000,1)

Sorry about this --- I hadn't pushed some changes. If you pull from master and try now, it should work.