From 5acabfc4c3ba5c4352e53d5b82205697324cb667 Mon Sep 17 00:00:00 2001
From: Jernej Sabadin <116955183+JSabadin@users.noreply.github.com>
Date: Tue, 3 Sep 2024 21:32:48 +0200
Subject: [PATCH] Config `pin_memory` Option (#67)

---
 luxonis_train/core/core.py    | 1 +
 luxonis_train/utils/config.py | 1 +
 2 files changed, 2 insertions(+)

diff --git a/luxonis_train/core/core.py b/luxonis_train/core/core.py
index 16953062..c683773c 100644
--- a/luxonis_train/core/core.py
+++ b/luxonis_train/core/core.py
@@ -167,6 +167,7 @@ def __init__(
                 drop_last=(
                     self.cfg.trainer.skip_last_batch if view == "train" else False
                 ),
+                pin_memory=self.cfg.trainer.pin_memory,
                 sampler=sampler if view == "train" else None,
             )
             for view in ["train", "val", "test"]
diff --git a/luxonis_train/utils/config.py b/luxonis_train/utils/config.py
index 44c00637..e3e4c0fb 100644
--- a/luxonis_train/utils/config.py
+++ b/luxonis_train/utils/config.py
@@ -283,6 +283,7 @@ class TrainerConfig(BaseModelExtraForbid):
     validation_interval: Literal[-1] | PositiveInt = 1
     num_log_images: NonNegativeInt = 4
     skip_last_batch: bool = True
+    pin_memory: bool = True
     log_sub_losses: bool = True
     save_top_k: Literal[-1] | NonNegativeInt = 3