diff --git a/air_llm/tests/test_compression.py b/air_llm/tests/test_compression.py
index bb8479a..bf7e3d9 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[k])
+                    self.assertAlmostEqual(aa[k], a_state_dict[k])
                 else:
-                    RMSE_loss = torch.sqrt(loss_fn(aa[k], a[k]))
+                    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