diff --git a/library/lumina_train_util.py b/library/lumina_train_util.py
index f224e86c..14a79bb2 100644
--- a/library/lumina_train_util.py
+++ b/library/lumina_train_util.py
@@ -688,9 +688,9 @@ def denoise(
                     noise_pred, dim=tuple(range(1, len(noise_pred.shape))), keepdim=True
                 )
                 # Iterate through batch
-                for noise_norm, max_new_norm, noise in zip(noise_norms, max_new_norms, noise_pred):
+                for i, (noise_norm, max_new_norm) in enumerate(zip(noise_norms, max_new_norms)):
                     if noise_norm >= max_new_norm:
-                        noise = noise * (max_new_norm / noise_norm)
+                        noise_pred[i] = noise_pred[i] * (max_new_norm / noise_norm)
         else:
             noise_pred = noise_pred_cond
 
diff --git a/lumina_train_network.py b/lumina_train_network.py
index 6b7e7d22..e1b45ac7 100644
--- a/lumina_train_network.py
+++ b/lumina_train_network.py
@@ -230,7 +230,7 @@ class LuminaNetworkTrainer(train_network.NetworkTrainer):
         self.noise_scheduler_copy = copy.deepcopy(noise_scheduler)
         return noise_scheduler
 
-    def encode_images_to_latents(self, args, accelerator, vae, images):
+    def encode_images_to_latents(self, args, vae, images):
         return vae.encode(images)
 
     # not sure, they use same flux vae