From b862400ea296072667cfeeaaaca4a1c25922d606 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Sun, 18 Aug 2024 05:11:00 +0300 Subject: [PATCH] IPEX fix AMP custom_fwd --- modules/intel/ipex/__init__.py | 2 ++ 1 file changed, 2 insertions(+) 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