From f94ef07460ec9693324ab861b4d50caff6bc353c Mon Sep 17 00:00:00 2001
From: Vladislav Kalinichenko
 <47711431+vladkalinichencko@users.noreply.github.com>
Date: Sat, 13 Jan 2024 18:36:30 +0300
Subject: [PATCH 1/5] Update device.py

add argument for index of preferred GPU
---
 nougat/utils/device.py | 33 ++++++++++++++++++++++-----------
 1 file changed, 22 insertions(+), 11 deletions(-)

diff --git a/nougat/utils/device.py b/nougat/utils/device.py
index 7748dc5..9825076 100644
--- a/nougat/utils/device.py
+++ b/nougat/utils/device.py
@@ -24,15 +24,26 @@ def default_batch_size():
         logging.warning("No GPU found. Conversion on CPU is very slow.")
     return batch_size
 
+def move_to_device(model, bf16: bool = True, cuda: bool = True, device_index: int = 0):
+   try:
+       if torch.backends.mps.is_available():
+           return model.to("mps")
+   except AttributeError:
+       pass
+   if bf16:
+       model = model.to(torch.bfloat16)
+   if cuda and torch.cuda.is_available():
+       model = model.to(f"cuda:{device_index}")
+   return model
 
-def move_to_device(model, bf16: bool = True, cuda: bool = True):
-    try:
-        if torch.backends.mps.is_available():
-            return model.to("mps")
-    except AttributeError:
-        pass
-    if bf16:
-        model = model.to(torch.bfloat16)
-    if cuda and torch.cuda.is_available():
-        model = model.to("cuda")
-    return model
+# def move_to_device(model, bf16: bool = True, cuda: bool = True):
+#     try:
+#         if torch.backends.mps.is_available():
+#             return model.to("mps")
+#     except AttributeError:
+#         pass
+#     if bf16:
+#         model = model.to(torch.bfloat16)
+#     if cuda and torch.cuda.is_available():
+#         model = model.to("cuda")
+#     return model

From 80c5980a16eab4819696ceb525471c6cb15546b5 Mon Sep 17 00:00:00 2001
From: Vladislav Kalinichenko
 <47711431+vladkalinichencko@users.noreply.github.com>
Date: Sat, 13 Jan 2024 18:47:43 +0300
Subject: [PATCH 2/5] Update predict.py

Added support for device index in command line arguments
---
 predict.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/predict.py b/predict.py
index b35ef9a..c80cc3b 100644
--- a/predict.py
+++ b/predict.py
@@ -82,6 +82,12 @@ def get_args():
         type=str,
         help="Provide page numbers like '1-4,7' for pages 1 through 4 and page 7. Only works for single PDF input.",
     )
+    parser.add_argument(
+        "--device-index",
+        type=int,
+        default=0,
+        help="Index of the preferred GPU device.",
+    )
     parser.add_argument("pdf", nargs="+", type=Path, help="PDF(s) to process.")
     args = parser.parse_args()
     if args.checkpoint is None or not args.checkpoint.exists():
@@ -125,7 +131,7 @@ def get_args():
 def main():
     args = get_args()
     model = NougatModel.from_pretrained(args.checkpoint)
-    model = move_to_device(model, bf16=not args.full_precision, cuda=args.batchsize > 0)
+    model = move_to_device(model, bf16=not args.full_precision, cuda=args.batchsize > 0, device_index=args.device_index)
     if args.batchsize <= 0:
         # set batch size to 1. Need to check if there are benefits for CPU conversion for >1
         args.batchsize = 1

From 975dfa1a952d382ac131fbe8ca766e11fa8fad4c Mon Sep 17 00:00:00 2001
From: Vladislav Kalinichenko <pinkod200586@gmail.com>
Date: Sat, 13 Jan 2024 19:43:35 +0300
Subject: [PATCH 3/5] removed comments

---
 nougat/utils/device.py | 12 ------------
 1 file changed, 12 deletions(-)

diff --git a/nougat/utils/device.py b/nougat/utils/device.py
index 9825076..ce7ba34 100644
--- a/nougat/utils/device.py
+++ b/nougat/utils/device.py
@@ -35,15 +35,3 @@ def move_to_device(model, bf16: bool = True, cuda: bool = True, device_index: in
    if cuda and torch.cuda.is_available():
        model = model.to(f"cuda:{device_index}")
    return model
-
-# def move_to_device(model, bf16: bool = True, cuda: bool = True):
-#     try:
-#         if torch.backends.mps.is_available():
-#             return model.to("mps")
-#     except AttributeError:
-#         pass
-#     if bf16:
-#         model = model.to(torch.bfloat16)
-#     if cuda and torch.cuda.is_available():
-#         model = model.to("cuda")
-#     return model

From 6047fc519350adc7ce32f2200a6f9777c20945ec Mon Sep 17 00:00:00 2001
From: Vladislav Kalinichenko <pinkod200586@gmail.com>
Date: Sat, 13 Jan 2024 20:19:57 +0300
Subject: [PATCH 4/5] Updated with the new argument

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 6aa70ed..0415e27 100644
--- a/README.md
+++ b/README.md
@@ -50,7 +50,7 @@ $ nougat path/to/directory -o output_directory
 
 ```
 usage: nougat [-h] [--batchsize BATCHSIZE] [--checkpoint CHECKPOINT] [--model MODEL] [--out OUT]
-              [--recompute] [--markdown] [--no-skipping] pdf [pdf ...]
+              [--recompute] [--markdown] [--no-skipping] [--device-index DEVICEINDEX] pdf [pdf ...]
 
 positional arguments:
   pdf                   PDF(s) to process.

From 9dec28f76b280336860e95ed66607bee893f67a9 Mon Sep 17 00:00:00 2001
From: Vladislav Kalinichenko <pinkod200586@gmail.com>
Date: Sat, 13 Jan 2024 21:39:40 +0300
Subject: [PATCH 5/5] removed strict=False arguments

---
 nougat/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nougat/model.py b/nougat/model.py
index d7e1321..bd3ebda 100644
--- a/nougat/model.py
+++ b/nougat/model.py
@@ -266,7 +266,7 @@ def __init__(
                     ]
                 else:
                     new_bart_state_dict[x] = bart_state_dict[x]
-            self.model.load_state_dict(new_bart_state_dict, strict=False)
+            self.model.load_state_dict(new_bart_state_dict)
 
     def add_special_tokens(self, list_of_tokens: List[str]):
         """