salilsdesai commited on
Commit
055b0e2
1 Parent(s): 6547184

Fix setting model device for CUDA

Browse files

Previously, trying to run the model when ```device``` is set to ```cuda:0``` gave the following error:

```
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument index in method wrapper_CUDA__index_select)
```

This was happening because ```device``` was only being used to set the input's device while the model's device was being left as CPU. The fix is to also set the model's device using ```device```

# Screenshots

Before Changes [CUDA] (broken):
![image.png](https://cdn-uploads.huggingface.co/production/uploads/658af1b52e1f5bf05c11ac77/by8L00Aqpa_4lPY-W3Yzh.png)

After Changes [CUDA] (fixed):
![image.png](https://cdn-uploads.huggingface.co/production/uploads/658af1b52e1f5bf05c11ac77/drin1qwBHnzhOElI6dBUT.png)

After Changes [CPU] (still works):
![image.png](https://cdn-uploads.huggingface.co/production/uploads/658af1b52e1f5bf05c11ac77/bqfyEvbCxNOUb0quodXb6.png)

Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -2,9 +2,11 @@ import gradio as gr
2
  from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
3
  import torch
4
 
5
- model = AutoModelForSeq2SeqLM.from_pretrained("Jayyydyyy/m2m100_418m_tokipona")
6
- tokenizer = AutoTokenizer.from_pretrained("facebook/m2m100_418M")
7
  device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 
 
 
8
  LANG_CODES = {
9
  "English":"en",
10
  "toki pona":"tl"
 
2
  from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
3
  import torch
4
 
 
 
5
  device = "cuda:0" if torch.cuda.is_available() else "cpu"
6
+
7
+ model = AutoModelForSeq2SeqLM.from_pretrained("Jayyydyyy/m2m100_418m_tokipona").to(device)
8
+ tokenizer = AutoTokenizer.from_pretrained("facebook/m2m100_418M")
9
+
10
  LANG_CODES = {
11
  "English":"en",
12
  "toki pona":"tl"