diff --git a/flux_train_network.py b/flux_train_network.py
index 5072c63d..001f7176 100644
--- a/flux_train_network.py
+++ b/flux_train_network.py
@@ -334,7 +334,7 @@ class FluxNetworkTrainer(train_network.NetworkTrainer):
         # If CDC is enabled, this will transform the noise with geometry-aware covariance
         noisy_model_input, timesteps, sigmas = flux_train_utils.get_noisy_model_input_and_timestep(
             args, noise_scheduler, latents, noise, accelerator.device, weight_dtype,
-            gamma_b_dataset=gamma_b_dataset, latents_npz_paths=latents_npz_paths, timestep_index=timestep_index
+            gamma_b_dataset=gamma_b_dataset, latents_npz_paths=latents_npz_paths
         )
 
         # pack latents and get img_ids