support torch GPU, only autoinit cuda in the buffer

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)

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