From d62dea9a2a9cd567d46d40a1e38af877e2b03b4d Mon Sep 17 00:00:00 2001
From: Kai Song <50285351+USTCKAY@users.noreply.github.com>
Date: Fri, 18 Aug 2023 12:26:38 +0800
Subject: [PATCH] [Custom Device]add lr_ratio support of adamw for custom
 device (#56386)

* [Custom Device]add lr_ratio support of adamw for mlu

* fix

* fix
---
 python/paddle/optimizer/adamw.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/python/paddle/optimizer/adamw.py b/python/paddle/optimizer/adamw.py
index 4e3990952c6..6f61e5ddd5b 100644
--- a/python/paddle/optimizer/adamw.py
+++ b/python/paddle/optimizer/adamw.py
@@ -182,6 +182,8 @@ class AdamW(Optimizer):
             if (
                 not core.is_compiled_with_cuda()
                 and not core.is_compiled_with_xpu()
+                and paddle.device.get_device().split(":")[0]
+                not in paddle.device.get_all_custom_device_type()
             ):
                 raise NotImplementedError("'lr_ratio' is unimplemented in CPU.")
 
-- 
GitLab