support torch GPU, only autoinit cuda in the buffer
MODIFY
accel/cuda/ops_cuda.py
accel/cuda/ops_cuda.py
# 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)

MODIFY
tinygrad/ops/ops_torch.py
tinygrad/ops/ops_torch.py
+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

support torch GPU, only autoinit cuda in the buffer