From 0b0cffe0e6e6030c796f4342d7b7603cb9f56dce Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 19 Mar 2025 02:29:32 -0700 Subject: [PATCH] Update vision.py --- unsloth/models/vision.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 699a26a051..c5cd57de97 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -177,7 +177,11 @@ def unsloth_base_fast_generate( if "attention_mask" in kwargs: torch._dynamo.mark_static(kwargs["attention_mask"], 0) torch._dynamo.mark_dynamic(kwargs["attention_mask"], 1) - pass + if "pixel_values" in kwargs: + print(kwargs["pixel_values"].shape) + if "token_type_ids" in kwargs: + torch._dynamo.mark_static(kwargs["token_type_ids"], 0) + torch._dynamo.mark_dynamic(kwargs["token_type_ids"], 1) # Fix generation_config cache_implementation = getattr(self.config, "cache_implementation", "static")