fp16

2024-06-16 17:54:41 +05:00
parent 9fb37f950a
commit 48876a130b
1 changed files with 6 additions and 1 deletions
@@ -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))