Fix tests

Xilinx · Jan 14, 2025 · b1b59a2 · b1b59a2
1 parent 18fff54
commit b1b59a2
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/src/brevitas/graph/equalize.py b/src/brevitas/graph/equalize.py
@@ -1384,7 +1384,7 @@ def _apply_rotate(
             # Only "weight" is rotated
             tensor_names_axis = [("weight", _get_input_axis(module))]
             # If rotations are fused or if the module is an orphan sink, transform is applied directly onto the tensor
-            rewriter_class = ModuleInstanceRegisterParametrization if insert_rotation_module or fuse_rotations else ModuleInstanceRegisterParametrization
+            rewriter_class = ModuleInstanceTransformTensor if insert_rotation_module or fuse_rotations else ModuleInstanceRegisterParametrization
             # Obtain rewriters for applying the rotations
             for tensor_name, axis in tensor_names_axis:
                 rewriter = rewriter_class(

diff --git a/tests/brevitas/graph/test_transforms.py b/tests/brevitas/graph/test_transforms.py
@@ -359,7 +359,8 @@ def forward(self, x):
     model_unfused.linear.weight.data = model_fused.linear.weight.data
 
     model_fused = ModuleInstanceTransformTensor(
-        model_fused.linear, "weight", rot_mat, rot_func, None, axis).apply(model_fused)
+        model_fused.linear, "weight", RotationWeightParametrization(rot_mat, rot_func, axis,
+                                                                    None)).apply(model_fused)
     model_unfused = ModuleInstanceRegisterParametrization(
         model_unfused.linear,
         "weight",