Skip to content

GitLab

Commit bcc9e93e authored by thomwolf's avatar thomwolf
Browse files

fix test

parent a3274ac4
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 2 deletions
tests/modeling_gpt2_test.py
View file @ bcc9e93e
...@@ -152,9 +152,10 @@ class GPT2ModelTest(unittest.TestCase): ...@@ -152,9 +152,10 @@ class GPT2ModelTest(unittest.TestCase):
self.parent.assertListEqual( self.parent.assertListEqual(
list(result["lm_logits"].size()), list(result["lm_logits"].size()),
[self.batch_size, self.n_choices, self.seq_length, total_voc]) [self.batch_size, self.n_choices, self.seq_length, total_voc])
self.parent.assertEqual(self.n_layer, len(result["presents"]))
self.parent.assertListEqual( self.parent.assertListEqual(
list(result["presents"].size()), list(result["presents"][0].size()),
[self.batch_size, self.n_choices, self.seq_length, total_voc]) [2, self.batch_size * self.n_choices, self.n_head, self.seq_length, self.n_embd // self.n_head])
def check_gpt2_lm_head_loss_output(self, result): def check_gpt2_lm_head_loss_output(self, result):
self.parent.assertListEqual( self.parent.assertListEqual(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment