From b14a76c695f0c30f51cf897830ae1f599c77d478 Mon Sep 17 00:00:00 2001
From: Martin Kozlovsky <martin.kozlovsky@luxonis.com>
Date: Thu, 23 Jan 2025 00:37:36 -0500
Subject: [PATCH] fix for rectangular images

---
 luxonis_train/nodes/heads/ghostfacenet_head.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/luxonis_train/nodes/heads/ghostfacenet_head.py b/luxonis_train/nodes/heads/ghostfacenet_head.py
index 8d9b66f4..56afb889 100644
--- a/luxonis_train/nodes/heads/ghostfacenet_head.py
+++ b/luxonis_train/nodes/heads/ghostfacenet_head.py
@@ -35,15 +35,16 @@ def __init__(
         """
         super().__init__(**kwargs)
         self.embedding_size = embedding_size
-        image_size = self.original_in_shape[1]
+        _, H, W = self.original_in_shape
 
         self.head = nn.Sequential(
             ConvModule(
                 self.in_channels,
                 self.in_channels,
-                kernel_size=(image_size // 32)
-                if image_size % 32 == 0
-                else (image_size // 32 + 1),
+                kernel_size=(
+                    H // 32 if H % 32 == 0 else H // 32 + 1,
+                    W // 32 if W % 32 == 0 else W // 32 + 1,
+                ),
                 groups=self.in_channels,
                 activation=False,
             ),