diff --git a/comfy_latent_upscaler.py b/comfy_latent_upscaler.py
index 047542b..852aa62 100644
--- a/comfy_latent_upscaler.py
+++ b/comfy_latent_upscaler.py
@@ -90,6 +90,10 @@ class LatentUpscaler:
 		lt = samples["samples"]
 		lt = model(lt)
 		del model
+		if "noise_mask" in samples.keys():
+			# expand the noise mask to the same shape as the latent
+			mask = torch.nn.functional.interpolate(samples['noise_mask'], scale_factor=float(scale_factor), mode='bicubic') 
+			return ({"samples": lt, "noise_mask": mask},)
 		return ({"samples": lt},)
 
 NODE_CLASS_MAPPINGS = {