Cleanup

Xilinx · Oct 29, 2024 · cb691c4 · cb691c4
1 parent 87626d4
commit cb691c4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/brevitas/graph/equalize.py b/src/brevitas/graph/equalize.py
@@ -1289,6 +1289,7 @@ def _apply_rotate(model: nn.Module, regions: List[Region], full_rotation_method=
             try:
                 # Build hadamard rotation matrix
                 rot_mat, K = get_hadK(hidden_dim)
+                rot_func = _apply_had_device
             except AssertionError as e:
                 print(f"Incomptible shapes {hidden_dim}")
                 if not insert_rotation_module:
@@ -1298,7 +1299,6 @@ def _apply_rotate(model: nn.Module, regions: List[Region], full_rotation_method=
                     rot_func = _apply_ort_device
                 print("Skipping layers")
                 continue
-            rot_func = _apply_had_device
 
         for name, indexes in region.srcs.items():
             module = region.get_module_from_name(name)