Layerwise fix

Xilinx · Jan 10, 2025 · 7d27050 · 7d27050
1 parent 8b71a53
commit 7d27050
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/brevitas/graph/base.py b/src/brevitas/graph/base.py
@@ -251,7 +251,7 @@ def apply(self, model: GraphModule) -> GraphModule:
                 weight = RotationWeightParametrization(
                     self.rot_mat, self.rot_func, self.axis, self.K)(weight)
                 # Modify the weights in-place
-                getattr(old_module, self.tensor_name).data = weight
+                setattr(old_module, self.tensor_name, torch.nn.Parameter(weight))
 
                 if hasattr(old_module, 'offload_params'):
                     old_module.offload_params(old_module)