mrfakename commited on
Commit
4db061f
1 Parent(s): 1329a39

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -0
app.py CHANGED
@@ -20,6 +20,9 @@ model_id = "yifeihu/TB-OCR-preview-0.1"
20
 
21
  DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
22
 
 
 
 
23
  model = AutoModelForCausalLM.from_pretrained(
24
  model_id,
25
  device_map=DEVICE,
 
20
 
21
  DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
22
 
23
+ if not torch.cuda.is_available():
24
+ ABOUT += "\n\n### This demo is running on CPU, it will be very slow..."
25
+
26
  model = AutoModelForCausalLM.from_pretrained(
27
  model_id,
28
  device_map=DEVICE,