Assign self.norm_type to input norm_type

facebookresearch · Feb 2, 2024 · 9a173bf · 9a173bf
1 parent 164cc0f
commit 9a173bf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/fairscale/nn/model_parallel/layers.py b/fairscale/nn/model_parallel/layers.py
@@ -178,7 +178,7 @@ def __init__(
         self.embedding_dim = embedding_dim
         self.padding_idx = padding_idx
         self.max_norm = max_norm
-        self.norm_type = scale_grad_by_freq
+        self.norm_type = norm_type
         self.scale_grad_by_freq = scale_grad_by_freq
         self.sparse = sparse
         self._weight = None