From b0eab338d1c495e0a2a235aa6a9df32a53f8ae94 Mon Sep 17 00:00:00 2001
From: "DESKTOP-UL3K6JV\\Daniq" <31.10.93@bk.ru>
Date: Sun, 16 Jun 2024 16:56:18 +0500
Subject: [PATCH] added movedim

---
 classifiers.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/classifiers.py b/classifiers.py
index 62346d8..e16e116 100644
--- a/classifiers.py
+++ b/classifiers.py
@@ -55,8 +55,10 @@ class YOLOClassify:
     CATEGORY = "SuperMasterBlasterLaser/ComfyUI_YOLO_Classifiers"
 
     def classify(self, yolo_classifier_model, image_to_classify):
+        image_to_classify = image_to_classify.movedim(3, 1)
+
         transform = transforms.Compose([
-            transforms.ToPILImage(mode='LA')
+            transforms.Resize(size=(640, 640))
         ])
 
         img = transform(image_to_classify).unsqueeze(0)