github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 15:08:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

shard to one device is to (#13519)

Browse Source 
* shard to one device is to

* fst
This commit is contained in:
George Hotz
2025-12-01 16:29:53 -08:00
committed by GitHub
parent 0b92fd30f5
commit 7307120311
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
tinygrad/tensor.py
View File

@@ -394,6 +394,7 @@ class Tensor(OpMixin):
```
"""
assert isinstance(self.device, str), "can't shard a MultiLazyBuffer"
if len(devices) == 1: return self.to(devices[0])
devices = tuple(canonicalize_device(x) for x in devices)
mlb = self.uop.shard(devices, self._resolve_dim(axis)) if axis is not None else self.uop.copy_to_device(devices)
return Tensor(mlb, device=devices, requires_grad=self.requires_grad)