fix: only import cupy if cuda available

Signed-off-by: Dev Singh <dev@devksingh.com>
This commit is contained in:
Dev Singh 2020-09-22 19:08:23 +00:00
parent ab40e66a2c
commit 9aa886293f

View File

@ -48,13 +48,16 @@ __all__ = [
]
import torch
import cupy as cp
import numpy as np
global device
device = "cuda:0" if torch.torch.cuda.is_available() else "cpu"
if device !== "cpu":
import cupy as cp
#todo: document completely
def set_device(self, new_device):