fix test extra issue (#3159)

This commit is contained in:
George Hotz
2024-01-17 11:58:08 -08:00
committed by GitHub
parent 9cc2577a08
commit ee83505fcc
4 changed files with 2 additions and 0 deletions

View File

@@ -24,12 +24,14 @@ class TestMultiTensor(unittest.TestCase):
X.to_((d0, d1))
for lb in X.lazydata.lbs:
assert lb.shape == (256,)
(X + X).realize()
def test_shard(self):
X = Tensor.ones(256).contiguous().realize()
X.shard_((d0, d1), 0)
for lb in X.lazydata.lbs:
assert lb.shape == (128,)
(X + X).realize()
def test_shard_same_device(self):
X = Tensor.ones(256).contiguous().realize()