no more toCPU path (#1624)

This commit is contained in:
George Hotz
2023-08-22 11:07:26 -07:00
committed by GitHub
parent 463dece63e
commit de1fcc418f
2 changed files with 5 additions and 4 deletions

View File

@@ -109,7 +109,7 @@ class Tensor:
return self
def detach(self): return Tensor(self.lazydata, device=self.device, requires_grad=False)
def numpy(self) -> np.ndarray: return self.lazydata.toCPU()
def numpy(self) -> np.ndarray: return self.to('CPU').lazydata.toCPU()
# TODO: if things are realized this won't work
def to_(self, device:str):
@@ -117,7 +117,7 @@ class Tensor:
self.lazydata.device = device
if self.grad: self.grad.to_(device)
def to(self, device:str):
def to(self, device:str) -> Tensor:
ret = Tensor(self.lazydata, device)
if self.grad: ret.grad = self.grad.to(device)
return ret