From 24dd0d52edfc32ab6f887f22752145255d8524dc Mon Sep 17 00:00:00 2001 From: wozeparrot Date: Wed, 30 Jul 2025 20:18:56 -0700 Subject: [PATCH] feat: test remove to cpu (#11444) --- tinygrad/tensor.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 53a7790a62..307c1ae17b 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -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: