From e193a92fe11a18fdfe3eedde129d38c776050632 Mon Sep 17 00:00:00 2001
From: Disty0 <disty@disty.xyz>
Date: Fri, 31 Jan 2025 18:46:07 +0300
Subject: [PATCH] IPEX fix torch.cuda.synchronize

---
 modules/intel/ipex/__init__.py  |  1 -
 modules/intel/ipex/attention.py |  1 -
 modules/intel/ipex/hijacks.py   | 26 +++++++++++++++++---------
 3 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/modules/intel/ipex/__init__.py b/modules/intel/ipex/__init__.py
index a765144b3..d00e79b24 100644
--- a/modules/intel/ipex/__init__.py
+++ b/modules/intel/ipex/__init__.py
@@ -37,7 +37,6 @@ def ipex_init(): # pylint: disable=too-many-statements
             torch.cuda.is_current_stream_capturing = lambda: False
             torch.cuda.set_device = torch.xpu.set_device
             torch.cuda.stream = torch.xpu.stream
-            torch.cuda.synchronize = torch.xpu.synchronize
             torch.cuda.Event = torch.xpu.Event
             torch.cuda.Stream = torch.xpu.Stream
             torch.Tensor.cuda = torch.Tensor.xpu
diff --git a/modules/intel/ipex/attention.py b/modules/intel/ipex/attention.py
index 42961377a..80a8cc3cc 100644
--- a/modules/intel/ipex/attention.py
+++ b/modules/intel/ipex/attention.py
@@ -1,5 +1,4 @@
 import os
-import math
 import torch
 from functools import cache, wraps
 
diff --git a/modules/intel/ipex/hijacks.py b/modules/intel/ipex/hijacks.py
index 4f3a03b4e..a1be40916 100644
--- a/modules/intel/ipex/hijacks.py
+++ b/modules/intel/ipex/hijacks.py
@@ -300,14 +300,6 @@ def torch_linspace(*args, device=None, **kwargs):
     else:
         return original_torch_linspace(*args, device=device, **kwargs)
 
-original_torch_Generator = torch.Generator
-@wraps(torch.Generator)
-def torch_Generator(device=None):
-    if check_device(device):
-        return original_torch_Generator(return_xpu(device))
-    else:
-        return original_torch_Generator(device)
-
 original_torch_load = torch.load
 @wraps(torch.load)
 def torch_load(f, map_location=None, *args, **kwargs):
@@ -318,6 +310,21 @@ def torch_load(f, map_location=None, *args, **kwargs):
     else:
         return original_torch_load(f, *args, map_location=map_location, **kwargs)
 
+original_torch_Generator = torch.Generator
+@wraps(torch.Generator)
+def torch_Generator(device=None):
+    if check_device(device):
+        return original_torch_Generator(return_xpu(device))
+    else:
+        return original_torch_Generator(device)
+
+@wraps(torch.cuda.synchronize)
+def torch_cuda_synchronize(device=None):
+    if check_device(device):
+        return torch.xpu.synchronize(return_xpu(device))
+    else:
+        return torch.xpu.synchronize(device)
+
 
 # Hijack Functions:
 def ipex_hijacks(legacy=True):
@@ -336,8 +343,9 @@ def ipex_hijacks(legacy=True):
     torch.zeros = torch_zeros
     torch.full = torch_full
     torch.linspace = torch_linspace
-    torch.Generator = torch_Generator
     torch.load = torch_load
+    torch.Generator = torch_Generator
+    torch.cuda.synchronize = torch_cuda_synchronize
 
     torch.backends.cuda.sdp_kernel = return_null_context
     torch.nn.DataParallel = DummyDataParallel