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()