Haystack: RAG Predict throwing an error when using a GPU

Created on 14 Nov 2020 · 3Comments · Source: deepset-ai/haystack

Describe the bug
I have no issue creating the generator, but creating it with use_gpu on will lead to an error when attempting to predict. Note that I have no issues with the FARMReader when using a gpu.

Error message

Inferencing Samples: 100%|██████████| 1/1 [00:00<00:00, 56.41 Batches/s]
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-17-42737cee0272> in <module>()
     10         question=question,
     11         documents=retriever_results,
---> 12         top_k=1
     13     )
     14 

/usr/local/lib/python3.6/dist-packages/haystack/generator/transformers.py in predict(self, question, documents, top_k)
    236         # Compute doc scores from docs_embedding
    237         doc_scores = torch.bmm(question_embedding.unsqueeze(1),
--> 238                                passage_embeddings.unsqueeze(0).transpose(1, 2)).squeeze(1)
    239 
    240         # TODO Need transformers 3.4.0

RuntimeError: Expected object of device type cuda but got device type cpu for argument #0 'result' in call to _th_bmm_out

Additional context
This is in google colab, I can't test this anywhere else as I don't have access to a GPU.
I did make sure to setup torch at the top:

!pip install urllib3==1.25.4 
!pip install torch==1.6.0+cu101 torchvision==0.6.1+cu101 -f https://download.pytorch.org/whl/torch_stable.html

Without this torch update, it doesn't throw the error, but it doesn't appear to be using the GPU as it greatly improves the speed of the DPR (Which also has use_gpu on).

To Reproduce
Should be able to reproduce by taking the RAG tutorial and flipping off to on for use_gpu

System: