Fix test.

This commit is contained in:
comfyanonymous 2026-01-10 01:32:28 -05:00 committed by GitHub
parent 9420c1bb2b
commit d4a54ff96c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -153,7 +153,7 @@ class TestMixedPrecisionOps(unittest.TestCase):
state_dict2 = model.state_dict()
# Verify layer1.weight is a QuantizedTensor with scale preserved
self.assertEqual(state_dict2["layer1.weight"], fp8_weight)
self.assertTrue(torch.equal(state_dict2["layer1.weight"], fp8_weight))
self.assertEqual(state_dict2["layer1.weight_scale"].item(), 3.0)
self.assertEqual(model.layer1.weight._layout_cls, "TensorCoreFP8E4M3Layout")