From 48876a130b2c241741f156187419c1e45e1328b7 Mon Sep 17 00:00:00 2001 From: "DESKTOP-UL3K6JV\\Daniq" <31.10.93@bk.ru> Date: Sun, 16 Jun 2024 17:54:41 +0500 Subject: [PATCH] fp16 --- classifiers.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) 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))