bug fix (#2426)

huggingface · Feb 9, 2024 · 9c1d5ba · 9c1d5ba
1 parent b0b867d
commit 9c1d5ba
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
@@ -1754,10 +1754,11 @@ def _prepare_ipex(self, *args):
         for obj in result:
             if isinstance(obj, torch.nn.Module):
                 model = obj
+                model.train()
             elif isinstance(obj, (torch.optim.Optimizer)):
                 optimizer = obj
         if optimizer is not None and model is not None:
-            dtype = torch.bfloat16 if self.state.mixed_precision == "bf16" else torch.float32
+            dtype = torch.bfloat16 if self.state.mixed_precision == "bf16" else None
             if self.device.type == "xpu" and is_xpu_available():
                 model = model.to(self.device)
                 model, optimizer = torch.xpu.optimize(