geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

feat: test remove to cpu (#11444)

Browse Source
This commit is contained in:
wozeparrot
2025-07-30 20:18:56 -07:00
committed by GitHub
parent c3cfcb50cb
commit 24dd0d52ed
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/tensor.py
+5 -2
View File
@@ -282,7 +282,7 @@ class Tensor(MathTrait):
# TODO: this is a hack for writing to DISK. remove with working assign
if isinstance(self.device, str) and self.device.startswith("DISK"):
if x.__class__ is not Tensor: x = Tensor(x, device="CPU", dtype=self.dtype)
cast(Buffer, self.contiguous().realize().uop.base.buffer).ensure_allocated().copyin(x._data())
self._buffer().copyin(x._data())
return self
if x.__class__ is not Tensor: x = Tensor(x, device=self.device, dtype=self.dtype)
if self.uop is x.uop: return self # a self assign is a NOOP
@@ -299,7 +299,10 @@ class Tensor(MathTrait):
"""
return Tensor(self.uop.detach(), device=self.device, requires_grad=False)
def _buffer(self) -> Buffer: return cast(Buffer, self.cast(self.dtype.base).contiguous().to("CPU").realize().uop.base.buffer)
def _buffer(self) -> Buffer:
x = self.cast(self.dtype.base).contiguous()
if isinstance(self.device, tuple): x = x.to("CPU")
return cast(Buffer, x.realize().uop.base.buffer).ensure_allocated()
def _data(self) -> memoryview: return self._buffer().as_buffer()
def data(self) -> memoryview: