notrain test
@@ -74,6 +74,16 @@ class TestOpt(unittest.TestCase):
assert len(CL.CACHE) == 10, "optimizer didn't fold conv-backward batchnorm"
Tensor.training = False
+ def test_fold_conv_batchnorm_notrain(self):
+ img = Tensor.ones(1,3,8,8)
+ c1 = nn.Conv2d(3,32,3)
+ bn = nn.BatchNorm2D(32, track_running_stats=False)
+ # precache the bn
+ img_conv = bn(c1(img)).relu().realize()
+ with CLCache():
+ img_conv = bn(c1(img)).relu().realize()
+ assert len(CL.CACHE) == 1, "optimizer didn't fold conv-batchnorm at test time"
+
def test_fold_conv_batchnorm(self):
Tensor.training = True
img = Tensor.ones(1,3,8,8)
