Fix flaky BetterTransformer test (#564)

* fix flaky test * fix flaky
huggingface · Dec 8, 2022 · 86375c4 · 86375c4
1 parent 1588a2e
commit 86375c4
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/tests/bettertransformer/test_bettertransformer_encoder.py b/tests/bettertransformer/test_bettertransformer_encoder.py
@@ -309,7 +309,10 @@ def get_batch(batch_size, avg_seqlen, max_sequence_length, seqlen_stdev, vocab_s
     mean_tensor = torch.Tensor([avg_seqlen]).expand(batch_size)
     stdev_tensor = torch.Tensor([seqlen_stdev]).expand(batch_size)
     lengths = torch.normal(mean_tensor, stdev_tensor).to(torch.int)
-    lengths = torch.clamp(lengths, min=0, max=max_sequence_length)
+
+    # need at least a sequence length of 1 for BetterTransformer to work
+    lengths = torch.clamp(lengths, min=1, max=max_sequence_length)
+
     tokens = torch.full(
         (batch_size, max_sequence_length),
         pad_idx,

diff --git a/tests/bettertransformer/testing_bettertransformer_utils.py b/tests/bettertransformer/testing_bettertransformer_utils.py
@@ -186,7 +186,10 @@ def get_batch(batch_size, avg_seqlen, max_sequence_length, seqlen_stdev, vocab_s
     mean_tensor = torch.Tensor([avg_seqlen]).expand(batch_size)
     stdev_tensor = torch.Tensor([seqlen_stdev]).expand(batch_size)
     lengths = torch.normal(mean_tensor, stdev_tensor).to(torch.int)
-    lengths = torch.clamp(lengths, min=0, max=max_sequence_length)
+
+    # need at least a sequence length of 1 for BetterTransformer to work
+    lengths = torch.clamp(lengths, min=1, max=max_sequence_length)
+
     tokens = torch.full(
         (batch_size, max_sequence_length),
         pad_idx,