diff --git a/jobs/process/GenerateProcess.py b/jobs/process/GenerateProcess.py
index 6f67a322..53b2d641 100644
--- a/jobs/process/GenerateProcess.py
+++ b/jobs/process/GenerateProcess.py
@@ -102,7 +102,7 @@ class GenerateProcess(BaseProcess):
             if self.model_config.is_lumina2:
                 arch = 'lumina2'
             sampler = get_sampler(
-                self.train_config.noise_scheduler,
+                self.generate_config.sampler,
                 {
                     "prediction_type": "v_prediction" if self.model_config.is_v_pred else "epsilon",
                 },