use CPU for test_schedule_ring (#9843)

* use CPU for test_schedule_ring

* why pre-commit is good
This commit is contained in:
chenyu
2025-04-10 23:20:53 -04:00
committed by GitHub
parent 7045920786
commit e0ec8be37d

View File

@@ -1,5 +1,5 @@
import unittest
from tinygrad import Tensor, Device
from tinygrad import Tensor
from tinygrad.helpers import Context
from tinygrad.ops import Ops
@@ -7,7 +7,7 @@ class TestRingAllReduce(unittest.TestCase):
def test_schedule_ring(self):
with Context(RING=2):
N = 6
ds = tuple(f"{Device.DEFAULT}:{i}" for i in range(N))
ds = tuple(f"CPU:{i}" for i in range(N))
t = Tensor.empty(N, N*100).shard(ds, axis=0).realize()
schedules = t.sum(0).schedule_with_vars()[0]
copies = [si for si in schedules if si.ast.op is Ops.COPY]