Merge pull request #60 from the-database/dev

fix cpu upscale pin memory error
the-database · Dec 19, 2024 · 5744bff · 5744bff
2 parents ec34d58 + 69505bc
commit 5744bff
Showing 1 changed file with 8 additions and 3 deletions.
diff --git a/MangaJaNaiConverterGui/backend/src/nodes/impl/pytorch/auto_split.py b/MangaJaNaiConverterGui/backend/src/nodes/impl/pytorch/auto_split.py
@@ -64,9 +64,14 @@ def _into_tensor(
             except Exception:
                 # Some arrays cannot be made writeable, and we need to copy them
                 img = np.copy(img)
-        input_tensor = (
-            torch.from_numpy(img).pin_memory().to(device, dtype, non_blocking=True)
-        )
+        if device == torch.device("cpu"):
+            input_tensor = (
+                torch.from_numpy(img).to(device, dtype, non_blocking=True)
+            )
+        else:
+            input_tensor = (
+                torch.from_numpy(img).pin_memory().to(device, dtype, non_blocking=True)
+            )
         return input_tensor
     finally:
         img.flags.writeable = writeable