19 lines
581 B
Python
19 lines
581 B
Python
import logging
|
|
import torch
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
|
|
def get_device(force_cpu):
|
|
# Check if CUDA can be used
|
|
if torch.cuda.is_available() and not force_cpu:
|
|
logging.info("CUDA detected. Running with GPU acceleration.")
|
|
device = torch.device("cuda")
|
|
elif force_cpu:
|
|
logging.info("CUDA detected, but overriding with option '--cpu'. Running with only CPU.")
|
|
device = torch.device("cpu")
|
|
else:
|
|
logging.info("CUDA is *NOT* detected. Running with only CPU.")
|
|
device = torch.device("cpu")
|
|
return device
|