[diffusion] fix: respect dit_precision config instead of hardcoded bfloat16 in DenoisingStagge (#24988)

Co-authored-by: niehen6174 <niehen.6174@gmail.com>
2026-05-15 13:22:24 +08:00
parent ad4994dc1d
commit 6cfa4c9c2f
1 changed files with 1 additions and 1 deletions
@@ -618,7 +618,7 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
        )

        # Setup precision and autocast settings
-        target_dtype = torch.bfloat16
+        target_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision]
        autocast_enabled = (
            target_dtype != torch.float32
        ) and not server_args.disable_autocast