diff --git a/comfy/utils.py b/comfy/utils.py
index 04926c1e4..985cd9a1b 100644
--- a/comfy/utils.py
+++ b/comfy/utils.py
@@ -853,19 +853,18 @@ def reshape_mask(input_mask, output_shape):
     dims = len(output_shape) - 2
 
     if dims == 1:
-        mask = input_mask
         scale_mode = "linear"
 
     if dims == 2:
-        mask = input_mask.reshape((-1, 1, input_mask.shape[-2], input_mask.shape[-1]))
+        input_mask = input_mask.reshape((-1, 1, input_mask.shape[-2], input_mask.shape[-1]))
         scale_mode = "bilinear"
 
     if dims == 3:
         if len(input_mask.shape) < 5:
-            mask = input_mask.reshape((1, 1, -1, input_mask.shape[-2], input_mask.shape[-1]))
+            input_mask = input_mask.reshape((1, 1, -1, input_mask.shape[-2], input_mask.shape[-1]))
         scale_mode = "trilinear"
 
-    mask = torch.nn.functional.interpolate(mask, size=output_shape[2:], mode=scale_mode)
+    mask = torch.nn.functional.interpolate(input_mask, size=output_shape[2:], mode=scale_mode)
     if mask.shape[1] < output_shape[1]:
         mask = mask.repeat((1, output_shape[1]) + (1,) * dims)[:,:output_shape[1]]
     mask = comfy.utils.repeat_to_batch_size(mask, output_shape[0])
diff --git a/comfy_extras/nodes_model_advanced.py b/comfy_extras/nodes_model_advanced.py
index ed14b61ac..f085bf12f 100644
--- a/comfy_extras/nodes_model_advanced.py
+++ b/comfy_extras/nodes_model_advanced.py
@@ -26,8 +26,8 @@ class X0(comfy.model_sampling.EPS):
 class ModelSamplingDiscreteDistilled(comfy.model_sampling.ModelSamplingDiscrete):
     original_timesteps = 50
 
-    def __init__(self, model_config=None):
-        super().__init__(model_config)
+    def __init__(self, model_config=None, zsnr=None):
+        super().__init__(model_config, zsnr=zsnr)
 
         self.skip_steps = self.num_timesteps // self.original_timesteps