From 7909786dbfdb9a7912d989c3ff0f970e4297df44 Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Fri, 28 Oct 2022 18:37:53 -0700
Subject: [PATCH] one more opt test

---
 test/external_test_opt.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/test/external_test_opt.py b/test/external_test_opt.py
index 16a43c898a..09cb266aac 100644
--- a/test/external_test_opt.py
+++ b/test/external_test_opt.py
@@ -59,6 +59,21 @@ class TestOpt(unittest.TestCase):
       assert len(CL.CACHE) == 3, "optimizer didn't fold batchnorm"
     Tensor.training = False
 
+  def test_fold_conv_sgd(self):
+    # TODO: with Tensor.training
+    Tensor.training = True
+    img = Tensor.ones(1,3,4,4)
+    c1 = nn.Conv2d(3,32,3)
+    opt = optim.SGD(optim.get_parameters(c1))
+    with CLCache():
+      opt.zero_grad()
+      c1(img).relu().sum().backward()
+      opt.step()
+      # TODO: this should be 4, but the sum output child stays around
+      # with pushing_permutes it can be 3
+      assert len(CL.CACHE) == 5, "optimizer didn't fold conv-backward SGD"
+    Tensor.training = False
+
   def test_fold_conv_batchnorm_sgd(self):
     # TODO: with Tensor.training
     Tensor.training = True