Text classification: RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

fastai 2.7.12
torch 2.0.0
cudatoolkit 11.7

The model is loaded to GPU with

fastai.torch_core.defaults.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
classifier = load_learner('/app/models/classifier.pkl')

However, when doing prediction on a string variable called text, receiving the following error:

    prediction = classifier.predict(text)
  File "/anaconda/envs/myenv/lib/python3.8/site-packages/fastai/learner.py", line 321, in predict
    inp,preds,_,dec_preds = self.get_preds(dl=dl, with_input=True, with_decoded=True)
  File "/anaconda/envs/myenv/lib/python3.8/site-packages/fastai/learner.py", line 315, in get_preds
    if reorder and hasattr(dl, 'get_idxs'): res = nested_reorder(res, tensor(idxs).argsort())
  File "/anaconda/envs/myenv/lib/python3.8/site-packages/fastai/torch_core.py", line 779, in nested_reorder
    elif is_listy(t): return type(t)(nested_reorder(t_, idxs) for t_ in t)
  File "/anaconda/envs/myenv/lib/python3.8/site-packages/fastai/torch_core.py", line 779, in <genexpr>
    elif is_listy(t): return type(t)(nested_reorder(t_, idxs) for t_ in t)
  File "/anaconda/envs/myenv/lib/python3.8/site-packages/fastai/torch_core.py", line 778, in nested_reorder
    if isinstance(t, (Tensor,L)): return t[idxs]
  File "/anaconda/envs/myenv/lib/python3.8/site-packages/fastai/torch_core.py", line 382, in __torch_function__
    res = super().__torch_function__(func, types, args, ifnone(kwargs, {}))
  File "/anaconda/envs/myenv/lib/python3.8/site-packages/torch/_tensor.py", line 1295, in __torch_function__
    ret = func(*args, **kwargs)
RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

Any thoughts?

Just taking a guess here…

Shouldn’t have to do this since it should default to the same device as the model.
Maybe try: learn.dls.device = ‘cuda’ ?

Thank you, but as you suspected, it’s already using cuda.

Adding

learn.dls.device = ‘cuda’

gives the same error.

Found a fix for this, but it’s weird.

Before calling fast.ai, spacy3 is used to predict a chunk of text. This chunk of text is then passed to the fast.ai classifier. Spacy model is also loaded on the GPU with spacy.require_gpu(). Forcing spacy back to cpu using spacy.require_cpu() before calling fast.ai model fixed the issue!

I have no idea how these two interfere and cause the error.

That makes sense…
load_learner() by default puts the model in the cpu mode.

You can change where the model/dl are processed using

model.to(‘cuda:0’)
dls.to(‘cuda:0’)

In your case,

classifier.to(‘cuda:0’)

may work as well…