Merge branch 'main' into workflow-cleanup

nod-ai · Dec 3, 2024 · bcafb2c · bcafb2c
2 parents 9c96fe9 + 544a2d8
commit bcafb2c
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/sharktank/tests/types/quantizers_test.py b/sharktank/tests/types/quantizers_test.py
@@ -75,7 +75,9 @@ def testPerAxisRoundtrip(self):
         torch.testing.assert_close(
             ssq.scale, torch.tensor([0.2, 0.4, 0.8], dtype=torch.float32)
         )
-        torch.testing.assert_close(ssq.reciprocal_scale, torch.tensor([5.0, 2.5, 1.25]))
+        torch.testing.assert_close(
+            ssq.reciprocal_scale, torch.tensor([5.0, 2.5, 1.25], dtype=torch.float32)
+        )
         self.assertIs(ssq.dtype, torch.float16)
 
     def testPerAxisQuantDequant(self):