From 1ba83cc7fa46ec7e7707c8e5662fa6b74652f62f Mon Sep 17 00:00:00 2001
From: qazal <77887910+Qazalin@users.noreply.github.com>
Date: Sun, 18 Aug 2024 23:35:42 +0800
Subject: [PATCH] split test_sgd_4convs_fuse [run_process_replay] (#6158)

---
 test/test_schedule.py | 33 ++++++++++++++++++++++-----------
 1 file changed, 22 insertions(+), 11 deletions(-)

diff --git a/test/test_schedule.py b/test/test_schedule.py
index 71cec09068..67ab3fdbbe 100644
--- a/test/test_schedule.py
+++ b/test/test_schedule.py
@@ -970,17 +970,16 @@ class TestSchedule(unittest.TestCase):
       check_schedule(opt.schedule_step(), 9)
 
   def test_sgd_4convs_fuse(self):
-    with Context(FUSE_CONV_BW=1):
-      with Tensor.train():
-        img = Tensor.empty(2,3,64,64)
-        c1 = nn.Conv2d(3,4,3,bias=False)
-        c2 = nn.Conv2d(4,8,3,bias=False)
-        c3 = nn.Conv2d(8,16,3,bias=False)
-        c4 = nn.Conv2d(16,32,3,bias=False)
-        opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4]))
-        opt.zero_grad()
-        c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward()
-        check_schedule(opt.schedule_step(), 19)
+    with Tensor.train():
+      img = Tensor.empty(2,3,64,64)
+      c1 = nn.Conv2d(3,4,3,bias=False)
+      c2 = nn.Conv2d(4,8,3,bias=False)
+      c3 = nn.Conv2d(8,16,3,bias=False)
+      c4 = nn.Conv2d(16,32,3,bias=False)
+      opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4]))
+      opt.zero_grad()
+      c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward()
+      check_schedule(opt.schedule_step(), 22)
 
   @unittest.skipUnless(is_dtype_supported(dtypes.half), "need half")
   def test_prefer_half_buffer(self):
@@ -1346,6 +1345,18 @@ class TestConvBW(unittest.TestCase):
     np.testing.assert_allclose(c1.weight.grad.numpy(), c1_torch.weight.grad.numpy(), atol=5e-4, rtol=1e-5)
     np.testing.assert_allclose(img.grad.numpy(), img_torch.grad.numpy(), atol=5e-4, rtol=1e-5)
 
+  def test_sgd_4convs_fuse(self):
+    with Tensor.train():
+      img = Tensor.empty(2,3,64,64)
+      c1 = nn.Conv2d(3,4,3,bias=False)
+      c2 = nn.Conv2d(4,8,3,bias=False)
+      c3 = nn.Conv2d(8,16,3,bias=False)
+      c4 = nn.Conv2d(16,32,3,bias=False)
+      opt = nn.optim.SGD(nn.state.get_parameters([c1, c2, c3, c4]))
+      opt.zero_grad()
+      c4(c3(c2(c1(img).relu()).relu()).relu()).relu().sum().backward()
+      self.check_schedule(opt.schedule_step(), 19)
+
 class TestIndexing(unittest.TestCase):
   def check_schedule(self, xt:Union[Tensor,List[Tensor]], cnt:int):
     with Context(FUSE_ARANGE=getenv("FUSE_ARANGE", 1)):