Update torch utils to handle passing in a device #82
Description
torch_device() in _torch_utils.py returns None if the device is not set to auto
should add an else statement to the function
Original:
def torch_device(device: str) -> str:
import torch
if device == "auto":
if torch.cuda.is_available():
return "cuda"
else:
return "cpu"
Improved:
def torch_device(device: str) -> str:
import torch
if device == "auto":
if torch.cuda.is_available():
return "cuda"
else:
return "cpu"
else:
return device