Update decoder.py

decoder.py

@@ -1,7 +1,8 @@
 import torch.nn as nn
 import torch
-from transformers import Qwen2_5_VLForConditionalGeneration, AutoConfig
- 
+from transformers import Qwen2_5_VLForConditionalGeneration, AutoConfig，BitsAndBytesConfig
+
+
 class SketchDecoder(nn.Module):
   """
   Autoregressive generative model 
@@ -23,10 +24,17 @@ class SketchDecoder(nn.Module):
           eos_token_id=self.eos_token_id,
           pad_token_id=self.pad_token_id)
 
+    quantization_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.float16,
+        bnb_4bit_quant_type="nf4"
+    )
+
     self.transformer = Qwen2_5_VLForConditionalGeneration.from_pretrained(
         "Qwen/Qwen2.5-VL-3B-Instruct", 
         config=config,
-        torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2",
+        quantization_config=quantization_config,
+        torch_dtype=torch.bfloat16, attn_implementation="sdpa",
         device_map ="cuda",
         ignore_mismatched_sizes=True
     )