update tests

Author: dimidagd

src/transformers/models/dinov3_vit/modular_dinov3_vit.py

@@ -39,7 +39,7 @@
 from ...pytorch_utils import compile_compatible_method_lru_cache
 from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
 from ...utils.backbone_utils import BackboneMixin
-from ...utils.generic import can_return_tuple, check_model_inputs
+from ...utils.generic import check_model_inputs
 from .configuration_dinov3_vit import DINOv3ViTConfig
 
 
@@ -555,4 +555,4 @@ def forward(
         )
 
 
-__all__ = ["DINOv3ViTModel", "DINOv3ViTPreTrainedModel", "DINOv3ViTBackbone", "DINOv3ViTForImageClassification"]
+__all__ = ["DINOv3ViTModel", "DINOv3ViTPreTrainedModel", "DINOv3ViTBackbone", "DINOv3ViTForImageClassification"]

tests/models/dinov3_vit/test_modeling_dinov3_vit.py

@@ -205,7 +205,9 @@ class Dinov3ModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
     attention_mask and seq_length.
     """
 
-    all_model_classes = (DINOv3ViTModel, DINOv3ViTBackbone DINOv3ViTForImageClassification) if is_torch_available() else ()
+    all_model_classes = (
+        (DINOv3ViTModel, DINOv3ViTBackbone, DINOv3ViTForImageClassification) if is_torch_available() else ()
+    )
     pipeline_model_mapping = (
         {
             "image-feature-extraction": DINOv3ViTModel,
@@ -277,13 +279,14 @@ def test_model_from_pretrained(self):
         model_name = "facebook/dinov3-vits16-pretrain-lvd1689m"
         model = DINOv3ViTModel.from_pretrained(model_name)
         self.assertIsNotNone(model)
-        
+
     @slow
     def test_model_for_image_classification_from_pretrained(self):
         model_name = "dimidagd/dinov3-vit7b16-pretrain-lvd1689m-imagenet1k-lc"
         model = DINOv3ViTForImageClassification.from_pretrained(model_name)
         self.assertIsNotNone(model)
 
+
 # We will verify our results on an image of cute cats
 def prepare_img():
     image = Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png")
@@ -301,12 +304,8 @@ def default_image_processor(self):
             else None
         )
 
-
     @slow
     def test_inference_lc_head_imagenet(self):
-#        tensor = torch.ones(1,3,224,224).to(model.device)
-#        expected_output_std = 0.7570638656616211
-#        expected_output_mean = 6.4013e-03
         model = DINOv3ViTModel.from_pretrained("facebook/dinov3-vits16-pretrain-lvd1689m").to(torch_device)
 
         image_processor = self.default_image_processor
@@ -316,11 +315,19 @@ def test_inference_lc_head_imagenet(self):
         # forward pass
         with torch.no_grad():
             outputs = model(**inputs)
+        predicted_class_idx = outputs.logits.argmax(-1).item()
+        # 283 is cat
+        self.assertEqual(predicted_class_idx, 283)
+
+        test_tensor = torch.ones(1, 3, 224, 224).to(model.device)
+        with torch.no_grad():
+            outputs = model(test_tensor)
 
-        
-#        self.assertAlmostEqual(outputs.logits.std().item(), expected_output_std, places=4)
-#        self.assertAlmostEqual(outputs.logits.mean().item(), expected_output_mean, places=4)
+        expected_output_std = 0.7570638656616211
+        expected_output_mean = 6.4013e-03
 
+        self.assertAlmostEqual(outputs.logits.std().item(), expected_output_std, places=4)
+        self.assertAlmostEqual(outputs.logits.mean().item(), expected_output_mean, places=4)
 
     @slow
     def test_inference_no_head(self):