This commit is contained in:
Yu Li
2023-12-01 16:40:54 -06:00
parent e4d50e0d8f
commit 7ac3aed8eb

View File

@@ -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().cpu().numpy()[0], 0.5)
self.assertLess(RMSE_loss.detach().cpu().item(), 0.5)