Cleanup

Xilinx · Oct 29, 2024 · ecf7b4f · ecf7b4f
1 parent 53bac08
commit ecf7b4f
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)