hongbo-miao/hongbomiao.com

machine-learning/convolutional-neural-network/src/utils/device.py

0 mins

Issues

import torch

if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cpu"