Revert the autocast modification.

Signed-off-by: yuanwu <yuan.wu@intel.com>
huggingface · Jan 10, 2025 · 31c09ef · 31c09ef
1 parent 9d95e55
commit 31c09ef
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/optimum/habana/diffusers/pipelines/pipeline_utils.py b/optimum/habana/diffusers/pipelines/pipeline_utils.py
@@ -159,7 +159,14 @@ def __init__(
                 )
 
             if self.gaudi_config.use_torch_autocast:
-                self.gaudi_config.declare_autocast_bf16_fp32_ops()
+                if bf16_full_eval:
+                    logger.warning(
+                        "`use_torch_autocast` is True in the given Gaudi configuration but "
+                        "`torch_dtype=torch.bfloat16` was given. Disabling mixed precision and continuing in bf16 only."
+                    )
+                    self.gaudi_config.use_torch_autocast = False
+                else:
+                    self.gaudi_config.declare_autocast_bf16_fp32_ops()
 
             if bf16_full_eval or self.gaudi_config.use_torch_autocast:
                 import diffusers