Refine xpu callback

openvinotoolkit · Nov 30, 2023 · 6d84a5d · 6d84a5d
1 parent 98ffd5d
commit 6d84a5d
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/src/otx/algorithms/anomaly/adapters/anomalib/callbacks/xpu.py b/src/otx/algorithms/anomaly/adapters/anomalib/callbacks/xpu.py
@@ -20,11 +20,10 @@ def __init__(self, device_idx=0):
 
     def on_fit_start(self, trainer, pl_module):
         """Applies IPEX optimization before training."""
-        if is_xpu_available():
-            pl_module.to(self.device)
-            model, optimizer = torch.xpu.optimize(trainer.model, optimizer=trainer.optimizers[0], dtype=torch.float32)
-            trainer.optimizers = [optimizer]
-            trainer.model = model
+        pl_module.to(self.device)
+        model, optimizer = torch.xpu.optimize(trainer.model, optimizer=trainer.optimizers[0])
+        trainer.optimizers = [optimizer]
+        trainer.model = model
 
     def on_train_batch_start(self, trainer, pl_module, batch, batch_idx):
         """Moves train batch tensors to XPU."""