diff --git a/modules/intel/ipex/__init__.py b/modules/intel/ipex/__init__.py
index 0f62c0244..79cafd155 100644
--- a/modules/intel/ipex/__init__.py
+++ b/modules/intel/ipex/__init__.py
@@ -149,6 +149,8 @@ def ipex_init(): # pylint: disable=too-many-statements
 
             # AMP:
             if legacy:
+                torch.xpu.amp.custom_fwd = torch.amp.custom_fwd
+                torch.xpu.amp.custom_bwd = torch.amp.custom_bwd
                 torch.cuda.amp = torch.xpu.amp
                 torch.is_autocast_enabled = torch.xpu.is_autocast_xpu_enabled
                 torch.get_autocast_gpu_dtype = torch.xpu.get_autocast_xpu_dtype