Fix failing Hiera, SwiftFormer & LED Model integration tests (#43225)

Sai-Suraj-27 · vasqu · web-flow · commit e977446e6326 · 2026-01-13T16:40:06.000+01:00
* Fix failing Hiera, LED ModelIntegrationTests

* swiftformer

* update expectations to a10

---------

Co-authored-by: vasqu &lt;antonprogamer@gmail.com&gt;
diff --git a/tests/models/hiera/test_modeling_hiera.py b/tests/models/hiera/test_modeling_hiera.py
@@ -557,7 +557,7 @@ def test_inference_image_classification_head(self):
 
         expected_slice = torch.tensor([[0.8028, 0.2409, -0.2254, -0.3712, -0.2848]]).to(torch_device)
 
-        torch.testing.assert_close(outputs.logits[0, :5], expected_slice, rtol=1e-4, atol=1e-4)
+        torch.testing.assert_close(outputs.logits[:, :5], expected_slice, rtol=1e-4, atol=1e-4)
 
     def test_inference_interpolate_pos_encoding(self):
         model = HieraModel.from_pretrained("facebook/hiera-tiny-224-hf").to(torch_device)
diff --git a/tests/models/led/test_modeling_led.py b/tests/models/led/test_modeling_led.py
@@ -537,7 +537,7 @@ def test_inference_head(self):
         expected_slice = torch.tensor(
             [[33.6507, 6.4572, 16.8089], [5.8739, -2.4238, 11.2902], [-3.2139, -4.3149, 4.2783]], device=torch_device
         )
-        torch.testing.assert_close(output[:, :3, :3], expected_slice, rtol=TOLERANCE, atol=TOLERANCE)
+        torch.testing.assert_close(output[0, :3, :3], expected_slice, rtol=TOLERANCE, atol=TOLERANCE)
 
     def test_seq_to_seq_generation(self):
         # this test requires 16GB of RAM
diff --git a/tests/models/swiftformer/test_modeling_swiftformer.py b/tests/models/swiftformer/test_modeling_swiftformer.py
@@ -260,5 +260,5 @@ def test_inference_image_classification_head(self):
         expected_shape = torch.Size((1, 1000))
         self.assertEqual(outputs.logits.shape, expected_shape)
 
-        expected_slice = torch.tensor([[-2.1703e00, 2.1107e00, -2.0811e00]]).to(torch_device)
-        torch.testing.assert_close(outputs.logits[0, :3], expected_slice, rtol=1e-4, atol=1e-4)
+        expected_slice = torch.tensor([[-2.1726, 2.1093, -2.0801]]).to(torch_device)
+        torch.testing.assert_close(outputs.logits[:, :3], expected_slice, rtol=1e-4, atol=1e-4)

Original file line number	Diff line number	Diff line change
`@@ -537,7 +537,7 @@ def test_inference_head(self):`
`537`	`537`	`expected_slice = torch.tensor(`
`538`	`538`	`[[33.6507, 6.4572, 16.8089], [5.8739, -2.4238, 11.2902], [-3.2139, -4.3149, 4.2783]], device=torch_device`
`539`	`539`	`)`
`540`		`- torch.testing.assert_close(output[:, :3, :3], expected_slice, rtol=TOLERANCE, atol=TOLERANCE)`
	`540`	`+ torch.testing.assert_close(output[0, :3, :3], expected_slice, rtol=TOLERANCE, atol=TOLERANCE)`
`541`	`541`
`542`	`542`	`def test_seq_to_seq_generation(self):`
`543`	`543`	`# this test requires 16GB of RAM`