Reduce dtype

pomonam · Mar 19, 2024 · 5aafa30 · 5aafa30
1 parent 5ccd596
commit 5aafa30
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/tests/gpu_tests/ddp_test.py b/tests/gpu_tests/ddp_test.py
@@ -213,15 +213,15 @@ def test_lr_pairwise_scores(self) -> None:
         new_pairwise_scores = self.analyzer.load_pairwise_scores(scores_name="ddp_qb")
 
         if LOCAL_RANK == 0:
-            print(f"Previous score: {pairwise_scores[ALL_MODULE_NAME][0]}")
+            print(f"Previous score: {pairwise_scores[ALL_MODULE_NAME][50]}")
             print(f"Previous shape: {pairwise_scores[ALL_MODULE_NAME].shape}")
-            print(f"New score: {new_pairwise_scores[ALL_MODULE_NAME][0]}")
+            print(f"New score: {new_pairwise_scores[ALL_MODULE_NAME][50]}")
             print(f"New shape: {new_pairwise_scores[ALL_MODULE_NAME].shape}")
             assert check_tensor_dict_equivalence(
                 pairwise_scores,
                 new_pairwise_scores,
-                atol=1e-5,
-                rtol=1e-3,
+                atol=1e-3,
+                rtol=1e-1,
             )
 
     @classmethod