diff --git a/air_llm/tests/test_compression.py b/air_llm/tests/test_compression.py
index 0b7334b..6ceca4a 100644
--- a/air_llm/tests/test_compression.py
+++ b/air_llm/tests/test_compression.py
@@ -37,4 +37,4 @@ class TestCompression(unittest.TestCase):
                     self.assertTrue(torch.equal(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
+                    self.assertLess(RMSE_loss.detach().cpu().numpy()[0], 0.5)
\ No newline at end of file