Hope this is better

bigscience-workshop · Jul 4, 2022 · cc465ca · cc465ca
1 parent 8ddfba9
commit cc465ca
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/evaluation/results/opt/convert_transformers_checkpoint_to_meg_ds.py b/evaluation/results/opt/convert_transformers_checkpoint_to_meg_ds.py
@@ -112,7 +112,7 @@ def merge_layers(layers, num_heads: int, hidden_size: int):
             return torch.reshape(
                 torch.cat(
                     [
-                        layer.view(num_heads, hidden_size // num_heads, 1)
+                        layer.view(num_heads, 1, hidden_size // num_heads)
                         for layer in layers
                     ],
                     dim=1
@@ -124,7 +124,7 @@ def merge_layers(layers, num_heads: int, hidden_size: int):
             return torch.reshape(
                 torch.cat(
                     [
-                        layer.view(num_heads, hidden_size // num_heads, 1, hidden_size)
+                        layer.view(num_heads, 1, hidden_size // num_heads, hidden_size)
                         for layer in layers
                     ],
                     dim=1