use vector norm

allenai · Nov 19, 2024 · 8592cee · 8592cee
1 parent 825eb26
commit 8592cee
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/olmo_core/nn/functional/__init__.py b/src/olmo_core/nn/functional/__init__.py
@@ -10,5 +10,5 @@
 
 
 def l2_normalize(x: torch.Tensor, dim: int = -1) -> torch.Tensor:
-    # NOTE: could also use F.normalize(), but that doesn't work with DTensor at the moment
-    return x / torch.linalg.norm(x, dim=dim, keepdim=True, dtype=torch.float32).type_as(x)
+    # NOTE: could also use F.normalize(), but that doesn't work with DTensor at the moment.
+    return x / torch.linalg.vector_norm(x, dim=dim, keepdim=True, dtype=torch.float32).type_as(x)