From 56df41ef2b5745d5b83e9fee0ca22e2a239f9be5 Mon Sep 17 00:00:00 2001
From: "hongbo.mo" <hongbo.mo@upai.com>
Date: Thu, 10 Apr 2025 16:46:11 +0800
Subject: [PATCH] bugfix: reshape mask to (B, H, W)

---
 nodes.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/nodes.py b/nodes.py
index 612452e..41cc56b 100644
--- a/nodes.py
+++ b/nodes.py
@@ -244,7 +244,7 @@ class ReduxAdvanced:
         for t in conditioning:
             n = [torch.cat((t[0], cond), dim=1), t[1].copy()]
             c.append(n)
-        return (c, image, masko)
+        return (c, image, masko.squeeze(-1))
 
 
 # A dictionary that contains all nodes you want to export with their names
@@ -258,4 +258,4 @@ NODE_CLASS_MAPPINGS = {
 NODE_DISPLAY_NAME_MAPPINGS = {
     "StyleModelApplySimple": "Apply style model (simple)",
     "ReduxAdvanced": "Apply Redux model (advanced)"
-}
\ No newline at end of file
+}