enable training mask2former and maskformer for transformers trainer (#…

…28277)

* fix get_num_masks output as [int] to int

* fix loss size from torch.Size([1]) to torch.Size([])

src/transformers/models/mask2former/modeling_mask2former.py

@@ -787,7 +787,7 @@ def get_num_masks(self, class_labels: torch.Tensor, device: torch.device) -> tor
         Computes the average number of target masks across the batch, for normalization purposes.
         """
         num_masks = sum([len(classes) for classes in class_labels])
-        num_masks_pt = torch.as_tensor([num_masks], dtype=torch.float, device=device)
+        num_masks_pt = torch.as_tensor(num_masks, dtype=torch.float, device=device)
         return num_masks_pt
 
 

src/transformers/models/maskformer/modeling_maskformer.py

@@ -1193,7 +1193,7 @@ def get_num_masks(self, class_labels: torch.Tensor, device: torch.device) -> tor
         Computes the average number of target masks across the batch, for normalization purposes.
         """
         num_masks = sum([len(classes) for classes in class_labels])
-        num_masks_pt = torch.as_tensor([num_masks], dtype=torch.float, device=device)
+        num_masks_pt = torch.as_tensor(num_masks, dtype=torch.float, device=device)
         return num_masks_pt
 
 

tests/models/mask2former/test_modeling_mask2former.py

@@ -190,7 +190,7 @@ def comm_check_on_output(result):
         comm_check_on_output(result)
 
         self.parent.assertTrue(result.loss is not None)
-        self.parent.assertEqual(result.loss.shape, torch.Size([1]))
+        self.parent.assertEqual(result.loss.shape, torch.Size([]))
 
 
 @require_torch

tests/models/maskformer/test_modeling_maskformer.py

@@ -189,7 +189,7 @@ def comm_check_on_output(result):
         comm_check_on_output(result)
 
         self.parent.assertTrue(result.loss is not None)
-        self.parent.assertEqual(result.loss.shape, torch.Size([1]))
+        self.parent.assertEqual(result.loss.shape, torch.Size([]))
 
 
 @require_torch