From a5c273aadb974f3d388d2c7911bb7d7f5123601d Mon Sep 17 00:00:00 2001
From: sko00o <lwylwy@outlook.com>
Date: Mon, 21 Apr 2025 16:37:21 +0800
Subject: [PATCH] fix: update mask dimension order in letterbox function

---
 nodes.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nodes.py b/nodes.py
index 405f96a..8febf1f 100644
--- a/nodes.py
+++ b/nodes.py
@@ -94,7 +94,7 @@ def letterbox(img, mask, w, h, desiredSize):
         mask=torch.nn.functional.interpolate(mask, size=(h,w), mode="bicubic")
         letterbox = torch.zeros(size=(b,1,desiredSize,desiredSize))
         letterbox[:, :, offsety:(offsety+h), offsetx:(offsetx+w)] += mask
-        mask = letterbox.view(b,1,desiredSize,desiredSize)
+        mask = letterbox.movedim(1, -1) # B C H W -> B H W C
     return (img, mask)
 
 def getBoundingBox(mask, w, h, relativeMargin, desiredSize):