fix head dim in rope (#140)

eole-nlp · Nov 1, 2024 · 39d0396 · 39d0396
1 parent 146c8f9
commit 39d0396
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/eole/modules/rope.py b/eole/modules/rope.py
@@ -35,7 +35,7 @@ def __init__(self, model_config):
         """
         super(RotaryPosition, self).__init__()
         self.model_config = model_config
-        self.dim_per_head = model_config.hidden_size // model_config.heads
+        self.dim_per_head = model_config.dim_per_head
         if model_config.rope_config.rotary_dim == 0:
             rotary_dim = self.dim_per_head
         else: