add 2 more custom devices intel_gpu and apple mps

2023-06-27 11:14:14 +08:00 · 2023-06-27 11:14:14 +08:00 · 3d0c0eb59d
parent 62c56ce27c
commit 3d0c0eb59d
1 changed files with 3 additions and 2 deletions
--- a/ppcls/static/train.py
+++ b/ppcls/static/train.py
@ -90,8 +90,9 @@ def main(args):
        fleet.init(is_collective=True)

    # assign the device
-    assert global_config[
-        "device"] in ["cpu", "gpu", "xpu", "npu", "mlu", "ascend"]
+    assert global_config["device"] in [
+        "cpu", "gpu", "xpu", "npu", "mlu", "ascend", "intel_gpu", "mps"
+    ]
    device = paddle.set_device(global_config["device"])

    # amp related config