re-run action

intel-analytics · Aug 9, 2022 · ab0c52e · ab0c52e
1 parent 08a86cf
commit ab0c52e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/nano/src/bigdl/nano/pytorch/strategies/ipex/ipex_strategy.py b/python/nano/src/bigdl/nano/pytorch/strategies/ipex/ipex_strategy.py
@@ -83,7 +83,8 @@ class IPEXBF16Precision(PrecisionPlugin):
     @contextmanager
     def forward_context(self):
         """AMP for managing model forward/training_step/evaluation_step/predict_step."""
-        with torch.cpu.amp.autocast():
+        # Manually set the dtype
+        with torch.cpu.amp.autocast(dtype=torch.bfloat16):
             yield
 
     def optimizer_step(self,