[CI] flaky unit test fix (#318)

facebookresearch · Jun 1, 2022 · 9e27c8f · 9e27c8f
1 parent cbf4526
commit 9e27c8f
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tests/test_global_attention.py b/tests/test_global_attention.py
@@ -11,6 +11,9 @@
 def test_global_attention():
     b, s, d = 2, 90, 40
 
+    torch.cuda.manual_seed(42)
+    torch.manual_seed(42)
+
     def test_ratio(global_attention_ratio: float):
         # Make sure that Global and Normal attention get the same results for the corresponding tokens
         a = torch.rand(b, s, d)