From a98190895ac19cb67492d68f0eacc0f760de3da0 Mon Sep 17 00:00:00 2001
From: Mark Bailey <mbbailey96@gmail.com>
Date: Sat, 26 Oct 2024 10:34:32 -0400
Subject: [PATCH] feat: use cuda

---
 src/main.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/main.py b/src/main.py
index ca56317..49954e7 100644
--- a/src/main.py
+++ b/src/main.py
@@ -26,7 +26,7 @@ def root():
         pretrained=True,
         assume_straight_pages=False,
         preserve_aspect_ratio=True,
-    )
+    ).cuda()
 
     result = {}
     i = 0
@@ -67,9 +67,9 @@ def file(file: Image):
             "db_resnet50",
             "vitstr_base",
             pretrained=True,
-            assume_straight_pages=False,
+            straighten_pages=True,
             preserve_aspect_ratio=True,
-        )
+        ).cuda()
         doc = DocumentFile.from_images("data/" + file_name)
         pred_res = predictor(doc)
         json_res = pred_res.export()