Fix tests

Xilinx · Jan 13, 2025 · c599f31 · c599f31
1 parent 82287b0
commit c599f31
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/brevitas/graph/equalize.py b/src/brevitas/graph/equalize.py
@@ -1351,8 +1351,9 @@ def _apply_rotate(
                 continue
 
         # Cast rotation matrix to the weight dtype
-        dtype = next(model.parameters()).dtype
-        rot_mat = rot_mat.to(dtype=dtype)
+        if rot_mat is not None:
+            dtype = next(model.parameters()).dtype
+            rot_mat = rot_mat.to(dtype=dtype)
         # If the rotation is not fused, redefine as a Parameter, to enable its optimization
         if not insert_rotation_module and not fuse_rotations:
             rot_mat = torch.nn.Parameter(rot_mat)