diff --git a/accel/cuda/ops_cuda.py b/accel/cuda/ops_cuda.py index ca95f971f4..50c992c689 100644 --- a/accel/cuda/ops_cuda.py +++ b/accel/cuda/ops_cuda.py @@ -1,11 +1,11 @@ # pip3 install pycuda import pycuda.driver as cuda -import pycuda.autoinit - import numpy as np class CudaBuffer: def __init__(self, shape, hostbuf=None): + import pycuda.autoinit + # TODO: these are generic self.shape = shape self.sz = int(np.prod(shape)*4) diff --git a/tinygrad/ops/ops_torch.py b/tinygrad/ops/ops_torch.py index 0cdb2d8a5d..c1b04d852a 100644 --- a/tinygrad/ops/ops_torch.py +++ b/tinygrad/ops/ops_torch.py @@ -1,15 +1,17 @@ +import os import torch import numpy as np from ..tensor import Function +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") class TorchBuffer(torch.Tensor): def custompad(x, padding): return torch.nn.functional.pad(x, [item for sublist in padding[::-1] for item in sublist]) @staticmethod def fromCPU(data): - return TorchBuffer(torch.from_numpy(data).requires_grad_(False)) + return TorchBuffer(torch.from_numpy(data).requires_grad_(False)).to(device) def toCPU(x): - return x.numpy() + return x.cpu().numpy() def getdtype(self): return np.float32