fix bugs

Signed-off-by: zhaohu xing <920232796@qq.com>
2022-12-06 09:03:55 +08:00
parent 9c86fb8cac
commit 4929503258
3 changed files with 70 additions and 3 deletions
@@ -38,8 +38,8 @@ def get_optimal_device():
    if torch.cuda.is_available():
        return torch.device(get_cuda_device_string())

-    # if has_mps():
-    #     return torch.device("mps")
+    if has_mps():
+        return torch.device("mps")

    return cpu