Fix tests

Xilinx · Jul 3, 2024 · 6f0d2f9 · 6f0d2f9
1 parent 4fd4998
commit 6f0d2f9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/brevitas/graph/test_transforms.py b/tests/brevitas/graph/test_transforms.py
@@ -287,6 +287,6 @@ def forward(self, x):
     model = TestModel()
     assert model.conv.stride == (1, 1)
 
-    kwargs = {'stride': lambda module: 2 if module.in_channels == 3 else 1}
+    kwargs = {'stride': lambda module, name: 2 if module.in_channels == 3 else 1}
     model = ModuleToModuleByInstance(model.conv, nn.Conv2d, **kwargs).apply(model)
     assert model.conv.stride == (2, 2)