MohamedRashad commited on
Commit
ba3e7a4
1 Parent(s): ea38d8e

Add device selection and print statement for device

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -7,7 +7,9 @@ from pdf2image import convert_from_path
7
  processor = NougatProcessor.from_pretrained("MohamedRashad/arabic-small-nougat")
8
  model = VisionEncoderDecoderModel.from_pretrained("MohamedRashad/arabic-small-nougat")
9
  device = "cuda" if torch.cuda.is_available() else "cpu"
 
10
 
 
11
  context_length = 2048
12
 
13
  def extract_text_from_image(image):
 
7
  processor = NougatProcessor.from_pretrained("MohamedRashad/arabic-small-nougat")
8
  model = VisionEncoderDecoderModel.from_pretrained("MohamedRashad/arabic-small-nougat")
9
  device = "cuda" if torch.cuda.is_available() else "cpu"
10
+ model.to(device)
11
 
12
+ print(f"Using {device} device")
13
  context_length = 2048
14
 
15
  def extract_text_from_image(image):