the-algorithm-ml/common/test_device.py
wiseaidev 2cc1abedd7 add & config mypy on common package
Signed-off-by: wiseaidev <business@wiseai.dev>
2023-04-01 14:38:59 +03:00

16 lines
349 B
Python

"""Minimal test for device.
Mostly a test that this can be imported properly even tho moved.
"""
from unittest.mock import (
patch,
)
import tml.common.device as device_utils
def test_device():
with patch("tml.common.device.dist.init_process_group"):
device = device_utils.setup_and_get_device(tf_ok=False)
assert device.type == "cpu"