def get_device()

in distilvit/train.py [0:0]

• 6 lines of code
• 3 McCabe index (conditional complexity)

def get_device():
    if torch.cuda.is_available():
        return "cuda"
    elif torch.backends.mps.is_available():
        return "mps"
    return "cpu"