laverdes commited on
Commit
ef55841
1 Parent(s): 43a5321

fix: add device and sent model to device

Browse files
Files changed (1) hide show
  1. app.py +4 -0
app.py CHANGED
@@ -69,6 +69,10 @@ st.image(image, caption='Your target receipt')
69
  st.text(f'baking the 🍩...')
70
  processor = DonutProcessor.from_pretrained("unstructuredio/donut-base-sroie")
71
  pretrained_model = VisionEncoderDecoderModel.from_pretrained("unstructuredio/donut-base-sroie")
 
 
 
 
72
  pretrained_model.encoder.to(torch.bfloat16)
73
  pretrained_model.eval()
74
 
 
69
  st.text(f'baking the 🍩...')
70
  processor = DonutProcessor.from_pretrained("unstructuredio/donut-base-sroie")
71
  pretrained_model = VisionEncoderDecoderModel.from_pretrained("unstructuredio/donut-base-sroie")
72
+
73
+ device = "cuda" if torch.cuda.is_available() else "cpu"
74
+ pretrained_model.to(device)
75
+
76
  pretrained_model.encoder.to(torch.bfloat16)
77
  pretrained_model.eval()
78