From 3198afc524fa4a8607462cdaf7f5f888df1fca63 Mon Sep 17 00:00:00 2001
From: Patrick Morgan <patrick@patrick-morgan.net>
Date: Sun, 15 Dec 2024 21:08:37 -0800
Subject: [PATCH] fix

---
 nodes/nodes.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/nodes/nodes.py b/nodes/nodes.py
index 10147b6..ac81180 100644
--- a/nodes/nodes.py
+++ b/nodes/nodes.py
@@ -390,9 +390,16 @@ class MaskContourProcessor:
         Returns:
             tuple: Single-element tuple containing the processed mask tensor
         """
+        # Ensure mask is a 3D tensor
+        if len(mask.shape) == 2:
+            mask = mask.unsqueeze(0)  # Add batch dimension if missing
+
         # Convert mask to numpy
         mask_np = mask.cpu().numpy() if isinstance(mask, torch.Tensor) else np.array(mask)
         
+        # Ensure mask is 2D by taking the first batch if needed
+        mask_np = mask_np[0]  # Take first batch
+        
         # Create output mask
         output_mask = mask_np.copy()