From bedc395df8ad9c0df105854ae82ce5f3a3f08237 Mon Sep 17 00:00:00 2001 From: Yu Li Date: Fri, 1 Dec 2023 16:31:43 -0600 Subject: [PATCH] fix ut --- air_llm/tests/test_compression.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/air_llm/tests/test_compression.py b/air_llm/tests/test_compression.py index bf7e3d9..3979a7f 100644 --- a/air_llm/tests/test_compression.py +++ b/air_llm/tests/test_compression.py @@ -34,7 +34,7 @@ class TestCompression(unittest.TestCase): for k in aa.keys(): if compression is None: - self.assertAlmostEqual(aa[k], a_state_dict[k]) + self.assertTrue(torch.eq(aa[k], a_state_dict[k])) else: RMSE_loss = torch.sqrt(loss_fn(aa[k], a_state_dict[k])) self.assertLess(RMSE_loss.detach().numpy()[0], 0.5) \ No newline at end of file