Update LoRa layer size

llama_diffusion_model.py

@@ -28,7 +28,7 @@ class CustomTransformerModel(PreTrainedModel):
 
     def __init__(self, config):
         super().__init__(config)
-        self.llama = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-3B", torch_dtype=torch.float16, device_map="auto", token=hf_token)
+        self.llama = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-3B-Instruct", torch_dtype=torch.float16, device_map="auto", token=hf_token)
         self.llama.resize_token_embeddings(config.vocab_size)
 
         for param in self.llama.parameters():
@@ -37,7 +37,7 @@ class CustomTransformerModel(PreTrainedModel):
             param.requires_grad = True
 
         lora_config = LoraConfig(
-            r=512, lora_alpha=512, lora_dropout=0.0,
+            r=64, lora_alpha=64, lora_dropout=0.0,
             target_modules=["q_proj", "v_proj", "k_proj", "o_proj"],
             bias="none", task_type=None
         )