From 5f359a87303b63e80f26ea5e030d164b3d006d8e Mon Sep 17 00:00:00 2001
From: alpertunga-bile <holy_celestine@hotmail.com>
Date: Mon, 26 Aug 2024 20:26:52 +0300
Subject: [PATCH] torch_dtype quantization is added for input gen

---
 image_caption.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/image_caption.py b/image_caption.py
index 1522e19..fefa489 100644
--- a/image_caption.py
+++ b/image_caption.py
@@ -148,7 +148,10 @@ class ImageCaptionNode:
 
         processor = transformers.AutoProcessor.from_pretrained(model_path)
 
-        inputs = processor(images=img, return_tensors="pt").to(dev)
+        try:
+            inputs = processor(images=img, return_tensors="pt").to(dev, req_torch_dtype)
+        except:
+            inputs = processor(images=img, return_tensors="pt").to(dev)
 
         out = model.generate(
             **inputs,