diff --git a/src/brevitas/graph/base.py b/src/brevitas/graph/base.py
index c3c4b82e9..4954e1c10 100644
--- 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)