Added squeeze

2024-06-16 17:14:25 +05:00
parent d9f85fcf1e
commit d449ff8f98
1 changed files with 1 additions and 1 deletions
@@ -56,7 +56,7 @@ class YOLOClassify:

    def classify(self, yolo_classifier_model, image_to_classify):
        image_to_classify = image_to_classify.permute(0, 3, 1, 2)
-
+        image_to_classify = image_to_classify.squeeze(0)
        transform = transforms.Compose([
            transforms.Resize(size=(640, 640))
        ])