Upload processor

processor.py

@@ -71,7 +71,7 @@ class CondViTProcessor(ImageProcessingMixin):
         BatchFeature
             pixel_values : torch.Tensor
                 Processed image tensor (B C H W)
-            category : torch.Tensor
+            category_indices : torch.Tensor
                 Categories indices (B)
         """
         use_cats = categories is not None
@@ -86,6 +86,8 @@ class CondViTProcessor(ImageProcessingMixin):
         data["pixel_values"] = torch.stack([self.process_img(img) for img in images])
 
         if use_cats:
-            data["category"] = torch.stack([self.process_cat(c) for c in categories])
+            data["category_indices"] = torch.stack(
+                [self.process_cat(c) for c in categories]
+            )
 
         return BatchFeature(data=data)