Fix test.

This commit is contained in:
comfyanonymous 2025-11-30 03:07:03 -05:00 committed by GitHub
parent e1e50e017c
commit b440a2bde9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -115,6 +115,7 @@ class TestMixedPrecisionOps(unittest.TestCase):
# Forward pass
input_tensor = torch.randn(5, 10, dtype=torch.bfloat16)
with torch.inference_mode():
output = model(input_tensor)
self.assertEqual(output.shape, (5, 40))