Parcourir la source

shutdown dht in any casy

justheuristic il y a 5 ans
Parent
commit
13234a2db1
1 fichiers modifiés avec 5 ajouts et 0 suppressions
  1. 5 0
      tests/test_moe.py

+ 5 - 0
tests/test_moe.py

@@ -66,3 +66,8 @@ def test_compute_expert_scores():
                     assert torch.allclose(logits[b, e], gx[b, ii[b][e]] + gy[b, jj[b][e]]), "compute_expert_scores returned incorrect score"
         finally:
             dht.shutdown()
+
+
+if __name__ == '__main__':
+    test_remote_module_call()
+    test_compute_expert_scores()