Add files via upload

This commit is contained in:
David Rotermund 2024-02-26 13:01:09 +01:00 committed by GitHub
parent 9fa12c258e
commit 4173e3306e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -0,0 +1,17 @@
import torch
import logging
def get_torch_device(mylogger: logging.Logger, force_to_cpu: bool) -> torch.device:
if torch.cuda.is_available():
device_name: str = "cuda:0"
else:
device_name = "cpu"
if force_to_cpu:
device_name = "cpu"
mylogger.info(f"Using device: {device_name}")
device: torch.device = torch.device(device_name)
return device