fix

mosaicml · Apr 9, 2024 · 473761f · 473761f
1 parent e3f942f
commit 473761f
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/tests/models/layers/test_dmoe.py b/tests/models/layers/test_dmoe.py
@@ -271,8 +271,4 @@ def test_fwd_equal_dmoe(seqlen: int, precision: str, mlp_type: str):
     with context:
         mpt_logits = mb_dmoe_model(token_ids).logits
         db_logits = torch_dmoe_model(token_ids).logits
-        assert mb_allclose(mpt_logits,
-                           db_logits,
-                           rtol=0.005,
-                           atol=0.01,
-                           verbose=True)
+        assert mb_allclose(mpt_logits, db_logits)