Highlights vector normalization

Author: PierreQuinton

src/torchjd/aggregation/config.py

@@ -70,10 +70,7 @@ def forward(self, matrix: Tensor) -> Tensor:
         units = torch.nan_to_num((matrix / (matrix.norm(dim=1)).unsqueeze(1)), 0.0)
         best_direction = torch.linalg.pinv(units) @ weights
 
-        if best_direction.norm() == 0:
-            unit_target_vector = torch.zeros_like(best_direction)
-        else:
-            unit_target_vector = best_direction / best_direction.norm()
+        unit_target_vector = torch.nn.functional.normalize(best_direction)
 
         length = torch.sum(matrix @ unit_target_vector)