updated vit

src/similarity/model_implements/vit_base.py

@@ -12,7 +12,12 @@ class VitBase():
     def extract_feature(self, imgs):
         features = []
         for img in imgs:
+            # Convert the image to RGB if it has 4 channels
+            if img.mode == 'RGBA':
+                img = img.convert('RGB')
+                
             inputs = self.feature_extractor(images=img, return_tensors="pt")
+            # print("input shape: ", inputs.shape)
             with torch.no_grad():
                 outputs = self.model(**inputs)
             last_hidden_states =  outputs.last_hidden_state