test

Signed-off-by: sam1373 <samuelkriman@gmail.com>
yzhang123 · May 1, 2023 · a373b37 · a373b37
1 parent e878b3d
commit a373b37
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nemo/collections/nlp/modules/common/megatron/attention.py b/nemo/collections/nlp/modules/common/megatron/attention.py
@@ -404,8 +404,8 @@ def forward(
         # =====================
 
         if self.transient_global_tokens:
-            hidden_states = _pad_to_multiple(hidden_states, self.global_tokens_spacing, 0)
-            avg_hidden_states = hidden_states.reshape(
+            avg_hidden_states = _pad_to_multiple(hidden_states, self.global_tokens_spacing, 0)
+            avg_hidden_states = avg_hidden_states.reshape(
                 self.global_tokens_spacing, -1, hidden_states.shape[-2], hidden_states.shape[-1]
             )
             avg_hidden_states = avg_hidden_states.mean(dim=0)