Fix tensor dtype mismatch (huggingface#36985)

* Fix tensor dtype mismatch

* update

* update

---------

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>

Author: 2 people

tests/models/mixtral/test_modeling_mixtral.py

@@ -486,7 +486,7 @@ def test_small_model_logits(self):
         # Note: Key 9 is currently set for MI300, but may need potential future adjustments for H100s,
         # considering differences in hardware processing and potential deviations in output.
         EXPECTED_LOGITS = {
-            7: torch.Tensor([[0.1670, 0.1620, 0.6094], [-0.8906, -0.1588, -0.6060], [0.1572, 0.1290, 0.7246]]).to(
+            7: torch.Tensor([[0.1640, 0.1621, 0.6093], [-0.8906, -0.1640, -0.6093], [0.1562, 0.1250, 0.7226]]).to(
                 torch_device
             ),
             8: torch.Tensor([[0.1631, 0.1621, 0.6094], [-0.8906, -0.1621, -0.6094], [0.1572, 0.1270, 0.7227]]).to(
@@ -499,6 +499,8 @@ def test_small_model_logits(self):
         with torch.no_grad():
             logits = model(dummy_input).logits
 
+        logits = logits.float()
+
         torch.testing.assert_close(
             logits[0, :3, :3], EXPECTED_LOGITS[self.cuda_compute_capability_major_version], atol=1e-3, rtol=1e-3
         )
@@ -525,7 +527,7 @@ def test_small_model_logits_batched(self):
         # considering differences in hardware processing and potential deviations in generated text.
         EXPECTED_LOGITS_LEFT = {
             7: torch.Tensor(
-                [[0.1750, 0.0537, 0.7007], [0.1750, 0.0537, 0.7007], [0.1750, 0.0537, 0.7007]],
+                [[0.1904, 0.0500, 0.7187], [0.1933, 0.0515, 0.7187], [0.2001, 0.0559, 0.7148]],
             ).to(torch_device),
             8: torch.Tensor([[0.1914, 0.0508, 0.7188], [0.1953, 0.0510, 0.7227], [0.1973, 0.0562, 0.7148]]).to(
                 torch_device
@@ -537,7 +539,7 @@ def test_small_model_logits_batched(self):
 
         EXPECTED_LOGITS_LEFT_UNPADDED = {
             7: torch.Tensor(
-                [[0.2212, 0.5200, -0.3816], [0.8213, -0.2313, 0.6069], [0.2664, -0.7090, 0.2468]],
+                [[0.2236, 0.5195, -0.3828], [0.8203, -0.2275, 0.6054], [0.2656, -0.7070, 0.2460]],
             ).to(torch_device),
             8: torch.Tensor([[0.2217, 0.5195, -0.3828], [0.8203, -0.2295, 0.6055], [0.2676, -0.7109, 0.2461]]).to(
                 torch_device
@@ -548,7 +550,7 @@ def test_small_model_logits_batched(self):
         }
 
         EXPECTED_LOGITS_RIGHT_UNPADDED = {
-            7: torch.Tensor([[0.2205, 0.1232, -0.1611], [-0.3484, 0.3030, -1.0312], [0.0742, 0.7930, 0.7969]]).to(
+            7: torch.Tensor([[0.2167, 0.1269, -0.1640], [-0.3496, 0.2988, -1.0312], [0.0688, 0.7929, 0.8007]]).to(
                 torch_device
             ),
             8: torch.Tensor([[0.2178, 0.1260, -0.1621], [-0.3496, 0.2988, -1.0312], [0.0693, 0.7930, 0.8008]]).to(
@@ -561,6 +563,7 @@ def test_small_model_logits_batched(self):
 
         with torch.no_grad():
             logits = model(dummy_input, attention_mask=attention_mask).logits
+        logits = logits.float()
 
         torch.testing.assert_close(
             logits[0, :3, :3], EXPECTED_LOGITS_LEFT[self.cuda_compute_capability_major_version], atol=1e-3, rtol=1e-3