fixed inplace init

Author: mkhona-nvidia

tests/normalized_optimizer_convergence_test.py

@@ -36,13 +36,9 @@ def __init__(self, input_size=784, hidden_size=128, num_classes=10, dim=0):
     def _initialize_oblique_weights(self, dim):
         """Initialize weights to be normalized for oblique optimization."""
         with torch.no_grad():
-            # Normalize  of oblique layers
-            self.fc1.weight.data = self.fc1.weight.data / self.fc1.weight.data.norm(dim=dim, keepdim=True).clamp(
-                min=1e-8
-            )
-            self.fc2.weight.data = self.fc2.weight.data / self.fc2.weight.data.norm(dim=dim, keepdim=True).clamp(
-                min=1e-8
-            )
+            # Normalize in-place for oblique layers
+            self.fc1.weight.data /= self.fc1.weight.data.norm(dim=dim, keepdim=True).clamp(min=1e-8)
+            self.fc2.weight.data /= self.fc2.weight.data.norm(dim=dim, keepdim=True).clamp(min=1e-8)
 
     def forward(self, x):
         x = x.view(x.size(0), -1)  # Flatten
@@ -196,7 +192,7 @@ def test_optimizer_modes_convergence(self, optimizer_class: torch.optim.Optimize
         # Re-initialize for row normalization
         with torch.no_grad():
             for param in model.get_oblique_parameters():
-                param.data = param.data / param.data.norm(dim=optimizer_kwargs["dim"], keepdim=True).clamp(min=1e-8)
+                param.data /= param.data.norm(dim=optimizer_kwargs["dim"], keepdim=True).clamp(min=1e-8)
 
         # Train model
         initial_loss, final_loss, final_accuracy = self._train_model(