diff --git a/py/city96_latent_upscaler.py b/py/city96_latent_upscaler.py
index 1d0b9f4..f1656e9 100644
--- a/py/city96_latent_upscaler.py
+++ b/py/city96_latent_upscaler.py
@@ -1,4 +1,5 @@
 # https://github.com/city96/SD-Latent-Upscaler
+# check if any issues on weights and updates and local install for same
 import torch
 import torch.nn as nn
 from safetensors.torch import load_file
@@ -79,4 +80,9 @@ class LatentUpscaler:
 		lt = samples["samples"]
 		lt = model(lt)
 		del model
-		return ({"samples": lt},)
\ No newline at end of file
+		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},)
+