compatible with MAC

Author: OuyangWenyu

tests/test_model_utils.py

@@ -0,0 +1,65 @@
+"""
+Author: Wenyu Ouyang
+Date: 2025-04-15 13:07:20
+LastEditTime: 2025-04-15 13:09:25
+LastEditors: Wenyu Ouyang
+Description:
+FilePath: /torchhydro/tests/test_model_utils.py
+Copyright (c) 2023-2024 Wenyu Ouyang. All rights reserved.
+"""
+
+import pytest
+import torch
+from torchhydro.models.model_utils import get_the_device
+
+
+@pytest.mark.parametrize(
+    "device_num, expected_device",
+    [
+        (-1, "cpu"),
+        ([-1], "cpu"),
+        (["-1"], "cpu"),
+    ],
+)
+def test_get_the_device_cpu(device_num, expected_device):
+    device = get_the_device(device_num)
+    assert device.type == expected_device
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is not available")
+@pytest.mark.parametrize(
+    "device_num, expected_device",
+    [
+        (0, "cuda"),
+        ([0], "cuda"),
+        (1, "cuda"),
+    ],
+)
+def test_get_the_device_cuda(device_num, expected_device):
+    device = get_the_device(device_num)
+    assert device.type == expected_device
+    assert device.index == (
+        device_num[0] if isinstance(device_num, list) else device_num
+    )
+
+
+@pytest.mark.skipif(
+    not hasattr(torch.backends, "mps") or not torch.backends.mps.is_available(),
+    reason="MPS is not available",
+)
+@pytest.mark.parametrize(
+    "device_num, expected_device",
+    [
+        (0, "mps"),
+        ([0], "mps"),
+        (1, "mps"),  # Should warn and default to mps:0
+    ],
+)
+def test_get_the_device_mps(device_num, expected_device):
+    if device_num != 0:
+        with pytest.warns(UserWarning, match="MPS only supports device 0"):
+            device = get_the_device(device_num)
+    else:
+        device = get_the_device(device_num)
+    assert device.type == expected_device
+    assert device.index == 0

torchhydro/models/model_utils.py

@@ -1,13 +1,14 @@
 """
 Author: Wenyu Ouyang
 Date: 2021-08-09 10:19:13
-LastEditTime: 2023-09-21 16:47:05
+LastEditTime: 2025-04-15 12:59:35
 LastEditors: Wenyu Ouyang
 Description: Some util functions for modeling
 FilePath: /torchhydro/torchhydro/models/model_utils.py
 Copyright (c) 2021-2022 Wenyu Ouyang. All rights reserved.
 """
 
+import contextlib
 from typing import Union
 import warnings
 import torch
@@ -20,7 +21,7 @@ def get_the_device(device_num: Union[list, int]):
     Parameters
     ----------
     device_num : Union[list, int]
-        number of the device -- -1 means "cpu" or 0, 1, ... means "cuda:x"
+        number of the device -- -1 means "cpu" or 0, 1, ... means "cuda:x" or "mps:x"
     """
     if device_num in [[-1], -1, ["-1"]]:
         return torch.device("cpu")
@@ -30,6 +31,16 @@ def get_the_device(device_num: Union[list, int]):
             if type(device_num) is not list
             else torch.device(f"cuda:{str(device_num[0])}")
         )
+    # Check for MPS (MacOS)
+    mps_available = False
+    with contextlib.suppress(AttributeError):
+        mps_available = torch.backends.mps.is_available()
+    if mps_available:
+        if device_num != 0:
+            warnings.warn(
+                f"MPS only supports device 0. Using 'mps:0' instead of {device_num}."
+            )
+        return torch.device("mps:0")
     if device_num not in [[-1], -1, ["-1"]]:
         warnings.warn("You don't have GPU, so have to choose cpu for models")
     return torch.device("cpu")