diff --git a/classifiers.py b/classifiers.py index b67f689..d3a6fe6 100644 --- a/classifiers.py +++ b/classifiers.py @@ -56,7 +56,12 @@ 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 / 255 + + if yolo_classifier_model.fp16: + image_to_classify = image_to_classify.half() + else: + image_to_classify = image_to_classify.float() + image_to_classify = image_to_classify.squeeze(0) transform = transforms.Compose([ transforms.Resize(size=(224, 224))