director/device.py

import torch


def pick_device():
    if torch.cuda.is_available():
        return "cuda", torch.float16

    if torch.backends.mps.is_available():
        return "mps", torch.float16

    return "cpu", torch.float32