Cast rotation matrix to the weight dtype

Xilinx · Jan 13, 2025 · 82287b0 · 82287b0
1 parent 32c9087
commit 82287b0
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/brevitas/graph/equalize.py b/src/brevitas/graph/equalize.py
@@ -1350,6 +1350,9 @@ def _apply_rotate(
                 print("Skipping layers")
                 continue
 
+        # Cast rotation matrix to the weight dtype
+        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)