mirror of
https://github.com/titanscouting/tra-analysis.git
synced 2024-12-27 01:59:08 +00:00
fix: only import cupy if cuda available
Signed-off-by: Dev Singh <dev@devksingh.com>
This commit is contained in:
parent
ab40e66a2c
commit
9aa886293f
@ -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):
|
||||
|
Loading…
Reference in New Issue
Block a user