Spaces:

dhanilka
/

CaptionGenerator

Runtime error

dhanilka commited on Oct 19, 2023

Commit

900c7ca

•

1 Parent(s): 7cfef6e

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -7,7 +7,7 @@ from PIL import Image
 model_id = "adept/fuyu-8b"
 dtype = torch.bfloat16
-device = "cuda"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = FuyuForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=dtype)

 model_id = "adept/fuyu-8b"
 dtype = torch.bfloat16
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = FuyuForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype=dtype)