From 14a93e79332f8a12453673a344d0dcd2e6d35f8c Mon Sep 17 00:00:00 2001
From: liuyuhui <liuyuhui@baidu.com>
Date: Thu, 29 Apr 2021 17:08:26 +0800
Subject: [PATCH] [Kunlun]add multi xpu support for PaddleClas about dygraph
 (#690)

* add multi xpu support for PaddleClas about dygraph

* add dygraph multi xpu support
---
 tools/train.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/tools/train.py b/tools/train.py
index c77c1cfc7..48e15676c 100644
--- a/tools/train.py
+++ b/tools/train.py
@@ -56,7 +56,16 @@ def main(args):
     config = get_config(args.config, overrides=args.override, show=True)
     # assign the place
     use_gpu = config.get("use_gpu", True)
-    place = paddle.set_device('gpu' if use_gpu else 'cpu')
+    use_xpu = config.get("use_xpu", False)
+    assert (
+        use_gpu and use_xpu
+    ) is not True, "gpu and xpu can not be true in the same time in static mode!"
+    if use_gpu:
+        place = paddle.set_device('gpu')
+    elif use_xpu:
+        place = paddle.set_device('xpu')
+    else:
+        place = paddle.set_device('cpu')
 
     trainer_num = paddle.distributed.get_world_size()
     use_data_parallel = trainer_num != 1