napatswift commited on
Commit
48e6d08
1 Parent(s): 3526a5a

Update device

Browse files
Files changed (1) hide show
  1. main.py +3 -0
main.py CHANGED
@@ -2,6 +2,9 @@ from mmocr.ocr import MMOCR
2
  import gradio as gr
3
  import cv2
4
  import sys
 
 
 
5
 
6
  ocr = MMOCR(det_config='model/det/config.py',
7
  det_ckpt='model/det/model.pth',
 
2
  import gradio as gr
3
  import cv2
4
  import sys
5
+ import torch
6
+
7
+ device = 'gpu' if torch.cuda.is_available() else 'cpu'
8
 
9
  ocr = MMOCR(det_config='model/det/config.py',
10
  det_ckpt='model/det/model.pth',