Merge pull request #13 from nostalgebraist/fix-test-v0.0.11

Correct shape in test_activations_dict_to_array input example
jalammar · Dec 26, 2020 · 8f96065 · 8f96065
2 parents 2ee44d3 + 90deb5c
commit 8f96065
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/lm_test.py b/tests/lm_test.py
@@ -27,8 +27,8 @@ def test_select_output_token_sample(self):
         assert result == torch.tensor(2)
 
     def test_activations_dict_to_array(self):
-        dict = {0:[[np.zeros((3,4))]],
-                1:[[np.zeros((3,4))]]}
+        dict = {0:[np.zeros((3,4))],
+                1:[np.zeros((3,4))]}
         activations = activations_dict_to_array(dict)
         assert activations.shape == (2,4,3)