diff --git a/Makefile b/Makefile
index e1c15c5..888056d 100644
--- a/Makefile
+++ b/Makefile
@@ -1,6 +1,6 @@
 .PHONY: quality style test docs
 
-check_dirs := src
+check_dirs := src examples
 
 # Check that source code meets quality standards
 
diff --git a/README.md b/README.md
index 0e1c270..1dfafca 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,5 @@
 # 🤗 PET
-Parameter-Efficient Tuning. Intergrated with 🤗 Accelerate to scale seamlessly to large models using PyTorch FSDP. 
+Parameter-Efficient Tuning methods enable . Intergrated with 🤗 Accelerate to scale seamlessly to large models using PyTorch FSDP. 
 
 Supported methods:
 
@@ -8,6 +8,87 @@ Supported methods:
 3. P-Tuning
 4. Prompt Tuning 
 
+## Getting started
+
+```python
+from transformers import AutoModelForSeq2SeqLM
+from pet import get_pet_config, get_pet_model
+model_name_or_path = "bigscience/mt0-large"
+tokenizer_name_or_path = "bigscience/mt0-large"
+
+config = {
+    "pet_type":"LORA",
+    "task_type":"SEQ_2_SEQ_LM",
+    "r": 8,
+    "lora_alpha": 32,
+    "lora_dropout": 0.1
+}
+pet_config = get_pet_config(config)
+
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
+model = get_pet_model(model, pet_config)
+model.print_trainable_parameters()
+# output: trainable params: 2359296 || all params: 1231940608 || trainable%: 0.19151053100118282
+```
+
+## PET + 🤗 Accelerate
+
+PET models work with 🤗 Accelerate out of the box. 
+For scaling to large models, you can leverage 🤗 Accelerate's PyTorch FSDP integration as shown below.
+PyTorch FSDP shards parameters, gradients and optimizer states across data parallel workers which enables
+large language models to fit on available hardware. 
+It also supports CPU offloading to further enable distributed training at scale. 
+
+```python
+from pet.utils.other import fsdp_auto_wrap_policy
+
+...
+
+if os.environ.get("ACCELERATE_USE_FSDP", None) is not None:
+    accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model)
+
+model = accelerator.prepare(model)
+```
+
+Example of parameter efficient tuning with `mt0-xxl` base model using 🤗 Accelerate is provided in `~examples/pet_lora_seq2seq_accelerate_fsdp.py`. 
+1. First run `accelerate config --config_file fsdp_config.yaml` and answer the questionaire. 
+Below are the contents of the config file.
+```
+command_file: null
+commands: null
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: FSDP
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config:
+  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
+  fsdp_backward_prefetch_policy: BACKWARD_PRE
+  fsdp_offload_params: true
+  fsdp_sharding_strategy: 1
+  fsdp_state_dict_type: FULL_STATE_DICT
+  fsdp_transformer_layer_cls_to_wrap: T5Block
+gpu_ids: null
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+megatron_lm_config: {}
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 2
+rdzv_backend: static
+same_network: true
+tpu_name: null
+tpu_zone: null
+use_cpu: false
+```
+2. run the below command to launch example script
+```
+accelerate launch --config_file fsdp_config.yaml examples/pet_lora_seq2seq_accelerate_fsdp.py
+```
+
+
 ## Models support matrix
 
 ### Sequence Classification
@@ -41,5 +122,7 @@ Supported methods:
 
 ## Caveats:
 1. Doesn't work currently with DeeSpeed ZeRO Stage-3. Extending support with DeeSpeed ZeRO Stage-3 is in backlog.
+2. When using `P_TUNING` or `PROMPT_TUNING` with `SEQ_2_SEQ` task, remember to remove the `num_virtual_token` virtual prompt predictions from the left side of the model outputs during evaluations. 
+
 
 
diff --git a/examples/pet_lora_seq2seq.ipynb b/examples/pet_lora_seq2seq.ipynb
new file mode 100644
index 0000000..f8d86c6
--- /dev/null
+++ b/examples/pet_lora_seq2seq.ipynb
@@ -0,0 +1,4815 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "5f93b7d1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import AutoModelForSeq2SeqLM\n",
+    "from pet import get_pet_config,get_pet_model, get_pet_model_state_dict\n",
+    "import torch\n",
+    "from datasets import load_dataset\n",
+    "import os\n",
+    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
+    "from transformers import AutoTokenizer\n",
+    "from torch.utils.data import DataLoader\n",
+    "from transformers import default_data_collator,get_linear_schedule_with_warmup\n",
+    "from tqdm import tqdm\n",
+    "from datasets import load_dataset\n",
+    "\n",
+    "device = \"cuda\"\n",
+    "model_name_or_path = \"bigscience/mt0-large\"\n",
+    "tokenizer_name_or_path = \"bigscience/mt0-large\"\n",
+    "\n",
+    "config = {\n",
+    "    \"pet_type\":\"LORA\",\n",
+    "    \"task_type\":\"SEQ_2_SEQ_LM\",\n",
+    "    \"r\":16,\n",
+    "    \"lora_alpha\": 32,\n",
+    "    \"lora_dropout\": 0.1\n",
+    "}\n",
+    "checkpoint_name = \"financial_sentiment_analysis_lora_v1.pt\"\n",
+    "text_column = \"sentence\"\n",
+    "label_column = \"text_label\"\n",
+    "max_length=128\n",
+    "lr = 1e-3\n",
+    "num_epochs = 3\n",
+    "batch_size=8\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "8d0850ac",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "trainable params: 4718592 || all params: 1234299904 || trainable%: 0.38228893842642636\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "PETModelForSeq2SeqLM(\n",
+       "  (base_model): LoRAModel(\n",
+       "    (model): MT5ForConditionalGeneration(\n",
+       "      (shared): Embedding(250112, 1024)\n",
+       "      (encoder): T5Stack(\n",
+       "        (embed_tokens): Embedding(250112, 1024)\n",
+       "        (block): ModuleList(\n",
+       "          (0): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (relative_attention_bias): Embedding(32, 16)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (1): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (2): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (3): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (4): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (5): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (6): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (7): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (8): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (9): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (10): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (11): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (12): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (13): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (14): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (15): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (16): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (17): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (18): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (19): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (20): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (21): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (22): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (23): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (final_layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "        (dropout): Dropout(p=0.1, inplace=False)\n",
+       "      )\n",
+       "      (decoder): T5Stack(\n",
+       "        (embed_tokens): Embedding(250112, 1024)\n",
+       "        (block): ModuleList(\n",
+       "          (0): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (relative_attention_bias): Embedding(32, 16)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (1): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (2): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (3): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (4): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (5): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (6): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (7): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (8): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (9): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (10): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (11): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (12): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (13): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (14): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (15): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (16): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (17): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (18): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (19): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (20): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (21): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (22): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "          (23): T5Block(\n",
+       "            (layer): ModuleList(\n",
+       "              (0): T5LayerSelfAttention(\n",
+       "                (SelfAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (1): T5LayerCrossAttention(\n",
+       "                (EncDecAttention): T5Attention(\n",
+       "                  (q): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                  (v): Linear(\n",
+       "                    in_features=1024, out_features=1024, bias=False\n",
+       "                    (lora_dropout): Dropout(p=0.1, inplace=False)\n",
+       "                    (lora_A): Linear(in_features=1024, out_features=16, bias=False)\n",
+       "                    (lora_B): Linear(in_features=16, out_features=1024, bias=False)\n",
+       "                  )\n",
+       "                  (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "              (2): T5LayerFF(\n",
+       "                (DenseReluDense): T5DenseGatedActDense(\n",
+       "                  (wi_0): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wi_1): Linear(in_features=1024, out_features=2816, bias=False)\n",
+       "                  (wo): Linear(in_features=2816, out_features=1024, bias=False)\n",
+       "                  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                  (act): NewGELUActivation()\n",
+       "                )\n",
+       "                (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "              )\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (final_layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "        (dropout): Dropout(p=0.1, inplace=False)\n",
+       "      )\n",
+       "      (lm_head): Linear(in_features=1024, out_features=250112, bias=False)\n",
+       "    )\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# creating model\n",
+    "pet_config = get_pet_config(config)\n",
+    "\n",
+    "model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)\n",
+    "model = get_pet_model(model, pet_config)\n",
+    "model.print_trainable_parameters()\n",
+    "model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "4ee2babf",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/sourab/miniconda3/envs/ml/lib/python3.10/site-packages/huggingface_hub/utils/_deprecation.py:97: FutureWarning: Deprecated argument(s) used in 'dataset_info': token. Will not be supported from version '0.12'.\n",
+      "  warnings.warn(message, FutureWarning)\n",
+      "Found cached dataset financial_phrasebank (/home/sourab/.cache/huggingface/datasets/financial_phrasebank/sentences_allagree/1.0.0/550bde12e6c30e2674da973a55f57edde5181d53f5a5a34c1531c53f93b7e141)\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6de075f8208349108291ac5ab7f5c980",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/1 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4b0e67b6d93f43e4b0f6a2f8978e4b0c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/3 [00:00<?, ?ba/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a9551029c9884529bda7421a99170b51",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/1 [00:00<?, ?ba/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'sentence': 'The order was valued at USD12 .2 m.',\n",
+       " 'label': 1,\n",
+       " 'text_label': 'neutral'}"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# loading dataset\n",
+    "dataset = load_dataset(\"financial_phrasebank\", 'sentences_allagree')\n",
+    "dataset = dataset[\"train\"].train_test_split(test_size=0.1)\n",
+    "dataset[\"validation\"] = dataset[\"test\"]\n",
+    "del(dataset[\"test\"])\n",
+    "\n",
+    "classes = dataset[\"train\"].features[\"label\"].names\n",
+    "dataset = dataset.map(\n",
+    "    lambda x: {\"text_label\": [classes[label] for label in x[\"label\"]]},\n",
+    "    batched=True,\n",
+    "    num_proc=1,\n",
+    "    \n",
+    ")\n",
+    "\n",
+    "dataset[\"train\"][0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "adf9608c",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4421971232434db1b6141e91fda2f6d7",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Running tokenizer on dataset:   0%|          | 0/3 [00:00<?, ?ba/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9b2ef793d93443949f4a5d5874d4bc05",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Running tokenizer on dataset:   0%|          | 0/1 [00:00<?, ?ba/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# data preprocessing\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n",
+    "def preprocess_function(examples):\n",
+    "    inputs = examples[text_column]\n",
+    "    targets = examples[label_column]\n",
+    "    model_inputs = tokenizer(inputs, max_length=max_length, padding=\"max_length\", truncation=True, return_tensors=\"pt\")\n",
+    "    labels = tokenizer(targets, max_length=3, padding=\"max_length\", truncation=True, return_tensors=\"pt\")\n",
+    "    labels = labels[\"input_ids\"]\n",
+    "    labels[labels==tokenizer.pad_token_id] = -100\n",
+    "    model_inputs[\"labels\"] = labels\n",
+    "    return model_inputs\n",
+    "\n",
+    "processed_datasets = dataset.map(\n",
+    "            preprocess_function,\n",
+    "            batched=True,\n",
+    "            num_proc=1,\n",
+    "            remove_columns=dataset[\"train\"].column_names,\n",
+    "            load_from_cache_file=False,\n",
+    "            desc=\"Running tokenizer on dataset\",\n",
+    "        )\n",
+    "\n",
+    "train_dataset = processed_datasets[\"train\"]\n",
+    "eval_dataset = processed_datasets[\"validation\"]\n",
+    "\n",
+    "train_dataloader = DataLoader(\n",
+    "        train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True\n",
+    "    )\n",
+    "eval_dataloader = DataLoader(eval_dataset, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True)\n",
+    "\n",
+    "\n",
+    "\n",
+    "    "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "f733a3c6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# optimizer and lr scheduler\n",
+    "optimizer = torch.optim.AdamW(model.parameters(), lr=lr)\n",
+    "lr_scheduler = get_linear_schedule_with_warmup(\n",
+    "    optimizer=optimizer,\n",
+    "    num_warmup_steps=0,\n",
+    "    num_training_steps=(len(train_dataloader) * num_epochs),\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "6b3a4090",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████████████████████████| 255/255 [00:53<00:00,  4.80it/s]\n",
+      "100%|███████████████████████████████████████████████████████████████| 29/29 [00:02<00:00, 14.16it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=0: train_ppl=tensor(13.6966, device='cuda:0') train_epoch_loss=tensor(2.6171, device='cuda:0') eval_ppl=tensor(1.0046, device='cuda:0') eval_epoch_loss=tensor(0.0046, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████████████████████████| 255/255 [00:52<00:00,  4.88it/s]\n",
+      "100%|███████████████████████████████████████████████████████████████| 29/29 [00:02<00:00, 14.20it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=1: train_ppl=tensor(1.5893, device='cuda:0') train_epoch_loss=tensor(0.4633, device='cuda:0') eval_ppl=tensor(1.0020, device='cuda:0') eval_epoch_loss=tensor(0.0020, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████████████████████████| 255/255 [00:52<00:00,  4.87it/s]\n",
+      "100%|███████████████████████████████████████████████████████████████| 29/29 [00:02<00:00, 14.18it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=2: train_ppl=tensor(1.3210, device='cuda:0') train_epoch_loss=tensor(0.2784, device='cuda:0') eval_ppl=tensor(1.0026, device='cuda:0') eval_epoch_loss=tensor(0.0026, device='cuda:0')\n"
+     ]
+    }
+   ],
+   "source": [
+    "# training and evaluation\n",
+    "model = model.to(device)\n",
+    "\n",
+    "for epoch in range(num_epochs):\n",
+    "    model.train()\n",
+    "    total_loss = 0\n",
+    "    for step, batch in enumerate(tqdm(train_dataloader)):\n",
+    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
+    "        outputs = model(**batch)\n",
+    "        loss = outputs.loss\n",
+    "        total_loss += loss.detach().float()\n",
+    "        loss.backward()\n",
+    "        optimizer.step()\n",
+    "        lr_scheduler.step()\n",
+    "        optimizer.zero_grad()\n",
+    "\n",
+    "    model.eval()\n",
+    "    eval_loss = 0\n",
+    "    eval_preds = []\n",
+    "    for step, batch in enumerate(tqdm(eval_dataloader)):\n",
+    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
+    "        with torch.no_grad():\n",
+    "            outputs = model(**batch)\n",
+    "        loss = outputs.loss\n",
+    "        eval_loss += loss.detach().float()\n",
+    "        eval_preds.extend(tokenizer.batch_decode(torch.argmax(outputs.logits, -1).detach().cpu().numpy(), skip_special_tokens=True))\n",
+    "\n",
+    "    eval_epoch_loss = eval_loss/len(train_dataloader)\n",
+    "    eval_ppl = torch.exp(eval_epoch_loss)\n",
+    "    train_epoch_loss = total_loss/len(eval_dataloader)\n",
+    "    train_ppl = torch.exp(train_epoch_loss)\n",
+    "    print(f\"{epoch=}: {train_ppl=} {train_epoch_loss=} {eval_ppl=} {eval_epoch_loss=}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "6cafa67b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "accuracy=98.23788546255507 % on the evaluation dataset\n",
+      "eval_preds[:10]=['neutral', 'neutral', 'positive', 'positive', 'neutral', 'neutral', 'neutral', 'neutral', 'neutral', 'neutral']\n",
+      "dataset['validation']['text_label'][:10]=['neutral', 'neutral', 'positive', 'positive', 'neutral', 'neutral', 'neutral', 'neutral', 'neutral', 'neutral']\n"
+     ]
+    }
+   ],
+   "source": [
+    "# print accuracy\n",
+    "correct =0\n",
+    "total = 0\n",
+    "for pred,true in zip(eval_preds, dataset[\"validation\"][\"text_label\"]):\n",
+    "    if pred.strip()==true.strip():\n",
+    "        correct+=1\n",
+    "    total+=1  \n",
+    "accuracy = correct/total*100\n",
+    "print(f\"{accuracy=} % on the evaluation dataset\")\n",
+    "print(f\"{eval_preds[:10]=}\")\n",
+    "print(f\"{dataset['validation']['text_label'][:10]=}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "a8de6005",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'base_model.model.encoder.block.0.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0011,  0.0126, -0.0392,  ...,  0.0117, -0.0080,  0.0227],\n",
+      "        [ 0.0239, -0.0152,  0.0174,  ..., -0.0209, -0.0093, -0.0250],\n",
+      "        [ 0.0197,  0.0023,  0.0045,  ..., -0.0073, -0.0106, -0.0038],\n",
+      "        ...,\n",
+      "        [-0.0116,  0.0015,  0.0146,  ...,  0.0017, -0.0237, -0.0077],\n",
+      "        [ 0.0143,  0.0330,  0.0059,  ...,  0.0124,  0.0045, -0.0180],\n",
+      "        [ 0.0095,  0.0024,  0.0352,  ...,  0.0064,  0.0230, -0.0452]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.0.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0161, -0.0037,  0.0014,  ...,  0.0039,  0.0235, -0.0148],\n",
+      "        [-0.0086, -0.0180,  0.0253,  ...,  0.0277,  0.0029, -0.0193],\n",
+      "        [ 0.0004,  0.0172,  0.0078,  ...,  0.0094, -0.0130,  0.0012],\n",
+      "        ...,\n",
+      "        [ 0.0010,  0.0024,  0.0202,  ...,  0.0259, -0.0117,  0.0107],\n",
+      "        [ 0.0256,  0.0210, -0.0293,  ..., -0.0210,  0.0059, -0.0250],\n",
+      "        [-0.0056, -0.0048,  0.0187,  ...,  0.0086, -0.0178,  0.0021]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.0.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0205, -0.0127,  0.0213,  ..., -0.0042,  0.0034,  0.0243],\n",
+      "        [ 0.0197, -0.0426, -0.0326,  ...,  0.0069, -0.0203,  0.0358],\n",
+      "        [-0.0216,  0.0100, -0.0084,  ...,  0.0446, -0.0540,  0.0452],\n",
+      "        ...,\n",
+      "        [ 0.0578, -0.0072,  0.0047,  ..., -0.0078, -0.0231,  0.0116],\n",
+      "        [ 0.0131,  0.0072,  0.0372,  ..., -0.0183, -0.0214, -0.0139],\n",
+      "        [ 0.0013,  0.0402, -0.0062,  ...,  0.0365,  0.0286, -0.0329]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.0.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0046, -0.0044,  0.0061,  ..., -0.0013, -0.0303, -0.0034],\n",
+      "        [ 0.0060,  0.0068, -0.0019,  ..., -0.0046, -0.0336, -0.0189],\n",
+      "        [-0.0039, -0.0053,  0.0089,  ..., -0.0128,  0.0059, -0.0216],\n",
+      "        ...,\n",
+      "        [ 0.0052, -0.0199, -0.0261,  ..., -0.0253,  0.0173, -0.0027],\n",
+      "        [ 0.0098, -0.0081, -0.0105,  ..., -0.0074, -0.0016, -0.0026],\n",
+      "        [ 0.0134, -0.0060,  0.0122,  ...,  0.0021, -0.0014,  0.0148]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.1.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0306, -0.0243, -0.0414,  ...,  0.0121,  0.0280,  0.0006],\n",
+      "        [-0.0425, -0.0292,  0.0076,  ...,  0.0322, -0.0244,  0.0346],\n",
+      "        [-0.0081,  0.0155, -0.0289,  ..., -0.0078,  0.0121, -0.0135],\n",
+      "        ...,\n",
+      "        [ 0.0161,  0.0303, -0.0498,  ..., -0.0148,  0.0022, -0.0115],\n",
+      "        [ 0.0205,  0.0197, -0.0203,  ..., -0.0424, -0.0252, -0.0418],\n",
+      "        [ 0.0098, -0.0085, -0.0391,  ...,  0.0195, -0.0180, -0.0565]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.1.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 1.3011e-02, -8.4894e-03,  1.4550e-02,  ..., -2.8281e-02,\n",
+      "         -1.5803e-02, -6.7014e-03],\n",
+      "        [-6.9723e-03,  2.5044e-02,  1.0645e-03,  ..., -1.9179e-02,\n",
+      "         -2.0764e-02, -1.5413e-02],\n",
+      "        [ 1.1783e-02, -1.5737e-02,  1.4881e-02,  ..., -3.9572e-03,\n",
+      "          1.2995e-02,  2.8500e-02],\n",
+      "        ...,\n",
+      "        [-1.0119e-02,  3.6680e-03,  2.4997e-03,  ..., -1.1717e-02,\n",
+      "          6.6043e-03, -3.2334e-03],\n",
+      "        [-1.1863e-02, -6.9381e-04, -1.1733e-02,  ..., -3.1885e-03,\n",
+      "         -1.4473e-02,  1.3692e-03],\n",
+      "        [-7.2537e-03,  2.4926e-02,  2.9916e-03,  ..., -6.9848e-04,\n",
+      "         -5.9038e-03, -5.7989e-05]], device='cuda:0'), 'base_model.model.encoder.block.1.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0199, -0.0045,  0.0448,  ..., -0.0153, -0.0625, -0.0284],\n",
+      "        [-0.0109, -0.0260, -0.0110,  ...,  0.0036,  0.0067, -0.0112],\n",
+      "        [ 0.0537,  0.0035, -0.0319,  ..., -0.0041, -0.0150,  0.0522],\n",
+      "        ...,\n",
+      "        [ 0.0215,  0.0431, -0.0165,  ...,  0.0086,  0.0072,  0.0050],\n",
+      "        [ 0.0327, -0.0188,  0.0122,  ...,  0.0043, -0.0070, -0.0060],\n",
+      "        [ 0.0123, -0.0429,  0.0311,  ..., -0.0028,  0.0158, -0.0010]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.1.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0024, -0.0146, -0.0314,  ...,  0.0111, -0.0105, -0.0042],\n",
+      "        [ 0.0165,  0.0012, -0.0041,  ...,  0.0039, -0.0018,  0.0084],\n",
+      "        [-0.0278,  0.0229, -0.0050,  ..., -0.0019, -0.0090,  0.0025],\n",
+      "        ...,\n",
+      "        [ 0.0228,  0.0125, -0.0167,  ..., -0.0163, -0.0046,  0.0134],\n",
+      "        [-0.0191, -0.0080,  0.0129,  ...,  0.0015, -0.0096, -0.0112],\n",
+      "        [ 0.0306, -0.0115, -0.0116,  ...,  0.0047, -0.0097, -0.0013]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.2.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0183, -0.0045, -0.0373,  ...,  0.0028, -0.0018,  0.0099],\n",
+      "        [-0.0376,  0.0284, -0.0107,  ..., -0.0196, -0.0107,  0.0114],\n",
+      "        [ 0.0004,  0.0118, -0.0360,  ..., -0.0594,  0.0541,  0.0094],\n",
+      "        ...,\n",
+      "        [ 0.0039, -0.0022,  0.0018,  ...,  0.0467,  0.0007, -0.0348],\n",
+      "        [-0.0036,  0.0106,  0.0272,  ...,  0.0413,  0.0270,  0.0300],\n",
+      "        [-0.0325, -0.0090,  0.0226,  ..., -0.0035,  0.0270, -0.0215]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.2.layer.0.SelfAttention.q.lora_B.weight': tensor([[-0.0052,  0.0233, -0.0136,  ..., -0.0086,  0.0216,  0.0165],\n",
+      "        [ 0.0174,  0.0217,  0.0273,  ..., -0.0157,  0.0139,  0.0167],\n",
+      "        [ 0.0038, -0.0342,  0.0208,  ...,  0.0044, -0.0041,  0.0206],\n",
+      "        ...,\n",
+      "        [-0.0028,  0.0279,  0.0068,  ..., -0.0180,  0.0110,  0.0245],\n",
+      "        [-0.0083,  0.0002,  0.0032,  ...,  0.0041,  0.0075,  0.0023],\n",
+      "        [ 0.0368, -0.0012, -0.0103,  ..., -0.0077, -0.0096, -0.0221]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.2.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0620, -0.0185, -0.0250,  ..., -0.0078,  0.0048,  0.0218],\n",
+      "        [ 0.0750, -0.0434,  0.0007,  ..., -0.0083,  0.0214,  0.0416],\n",
+      "        [ 0.0181, -0.0213,  0.0101,  ...,  0.0022,  0.0036,  0.0140],\n",
+      "        ...,\n",
+      "        [-0.0114,  0.0142, -0.0239,  ..., -0.0055,  0.0138,  0.0047],\n",
+      "        [-0.0246, -0.0030,  0.0016,  ...,  0.0156, -0.0082, -0.0074],\n",
+      "        [-0.0391,  0.0287, -0.0008,  ..., -0.0037, -0.0228, -0.0337]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.2.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0043,  0.0104,  0.0072,  ..., -0.0040,  0.0189, -0.0040],\n",
+      "        [ 0.0302,  0.0278,  0.0314,  ..., -0.0275,  0.0137, -0.0292],\n",
+      "        [-0.0450, -0.0463, -0.0416,  ...,  0.0435, -0.0480,  0.0439],\n",
+      "        ...,\n",
+      "        [-0.0043, -0.0020,  0.0152,  ..., -0.0201,  0.0059, -0.0025],\n",
+      "        [-0.0169, -0.0239, -0.0168,  ...,  0.0162,  0.0110,  0.0223],\n",
+      "        [ 0.0274,  0.0228,  0.0166,  ..., -0.0190, -0.0184, -0.0188]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.3.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0117, -0.0470, -0.0126,  ...,  0.0142, -0.0365, -0.0033],\n",
+      "        [ 0.0198,  0.0106,  0.0006,  ..., -0.0219,  0.0049,  0.0288],\n",
+      "        [ 0.0024,  0.0177,  0.0296,  ..., -0.0142,  0.0018, -0.0186],\n",
+      "        ...,\n",
+      "        [-0.0298, -0.0051,  0.0024,  ...,  0.0243, -0.0158,  0.0213],\n",
+      "        [-0.0088, -0.0265,  0.0339,  ..., -0.0116, -0.0214, -0.0225],\n",
+      "        [-0.0259, -0.0074,  0.0242,  ...,  0.0003, -0.0149,  0.0445]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.3.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0196,  0.0119, -0.0112,  ...,  0.0255, -0.0357, -0.0064],\n",
+      "        [-0.0150,  0.0146,  0.0108,  ..., -0.0005, -0.0003,  0.0082],\n",
+      "        [-0.0081, -0.0125,  0.0033,  ..., -0.0002,  0.0076, -0.0156],\n",
+      "        ...,\n",
+      "        [-0.0096,  0.0218, -0.0031,  ...,  0.0057, -0.0008,  0.0058],\n",
+      "        [ 0.0084, -0.0011,  0.0069,  ..., -0.0263,  0.0075, -0.0034],\n",
+      "        [ 0.0040,  0.0189,  0.0171,  ..., -0.0186,  0.0084,  0.0118]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.3.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0178,  0.0288,  0.0192,  ...,  0.0078, -0.0146, -0.0158],\n",
+      "        [-0.0173,  0.0058, -0.0253,  ..., -0.0172, -0.0134, -0.0071],\n",
+      "        [-0.0167, -0.0333, -0.0169,  ..., -0.0249, -0.0090,  0.0195],\n",
+      "        ...,\n",
+      "        [ 0.0093, -0.0083,  0.0031,  ...,  0.0143,  0.0326,  0.0302],\n",
+      "        [-0.0282,  0.0131,  0.0045,  ..., -0.0152,  0.0099, -0.0038],\n",
+      "        [-0.0027, -0.0236, -0.0395,  ..., -0.0288,  0.0060, -0.0122]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.3.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0041,  0.0012,  0.0171,  ...,  0.0103, -0.0087, -0.0037],\n",
+      "        [ 0.0061, -0.0059,  0.0004,  ...,  0.0228, -0.0394,  0.0151],\n",
+      "        [-0.0123, -0.0029,  0.0168,  ...,  0.0022, -0.0009,  0.0042],\n",
+      "        ...,\n",
+      "        [-0.0318,  0.0018,  0.0162,  ..., -0.0191,  0.0094,  0.0015],\n",
+      "        [-0.0176,  0.0092,  0.0165,  ..., -0.0086,  0.0150,  0.0036],\n",
+      "        [ 0.0359,  0.0136, -0.0198,  ...,  0.0134, -0.0206, -0.0085]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.4.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0325,  0.0371,  0.0161,  ..., -0.0039, -0.0107, -0.0066],\n",
+      "        [-0.0440, -0.0042,  0.0123,  ..., -0.0205, -0.0390,  0.0127],\n",
+      "        [ 0.0115, -0.0246, -0.0256,  ..., -0.0113,  0.0181, -0.0191],\n",
+      "        ...,\n",
+      "        [-0.0033,  0.0258,  0.0214,  ...,  0.0249, -0.0544,  0.0062],\n",
+      "        [-0.0294, -0.0302, -0.0316,  ..., -0.0032, -0.0195,  0.0247],\n",
+      "        [-0.0197, -0.0115,  0.0108,  ...,  0.0308, -0.0469, -0.0307]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.4.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0078, -0.0045,  0.0053,  ...,  0.0212,  0.0162, -0.0151],\n",
+      "        [-0.0228,  0.0114,  0.0248,  ...,  0.0080,  0.0110, -0.0130],\n",
+      "        [-0.0307,  0.0002,  0.0129,  ..., -0.0120, -0.0139, -0.0271],\n",
+      "        ...,\n",
+      "        [-0.0056, -0.0221,  0.0047,  ...,  0.0003,  0.0098, -0.0024],\n",
+      "        [ 0.0210,  0.0042,  0.0005,  ...,  0.0077, -0.0075,  0.0073],\n",
+      "        [ 0.0255,  0.0454,  0.0152,  ...,  0.0398,  0.0391, -0.0117]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.4.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0198, -0.0339,  0.0026,  ...,  0.0093,  0.0370,  0.0019],\n",
+      "        [-0.0242, -0.0109,  0.0368,  ..., -0.0427, -0.0158, -0.0527],\n",
+      "        [-0.0196,  0.0233,  0.0298,  ...,  0.0180, -0.0326,  0.0129],\n",
+      "        ...,\n",
+      "        [ 0.0229,  0.0284, -0.0155,  ...,  0.0192,  0.0025,  0.0028],\n",
+      "        [ 0.0174, -0.0308,  0.0193,  ...,  0.0195, -0.0141,  0.0278],\n",
+      "        [ 0.0147,  0.0028, -0.0325,  ..., -0.0281,  0.0010,  0.0087]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.4.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0090,  0.0005, -0.0142,  ..., -0.0017, -0.0026, -0.0015],\n",
+      "        [ 0.0117, -0.0037,  0.0098,  ..., -0.0013,  0.0009, -0.0046],\n",
+      "        [-0.0163, -0.0004, -0.0039,  ...,  0.0007,  0.0109,  0.0066],\n",
+      "        ...,\n",
+      "        [-0.0187,  0.0198,  0.0076,  ...,  0.0136,  0.0249, -0.0021],\n",
+      "        [ 0.0008,  0.0045, -0.0036,  ..., -0.0057,  0.0173,  0.0004],\n",
+      "        [-0.0046, -0.0078, -0.0104,  ..., -0.0017, -0.0123,  0.0134]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.5.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 4.6854e-02,  2.2295e-02,  6.4651e-03,  ...,  2.5293e-02,\n",
+      "          2.1630e-02,  1.2340e-02],\n",
+      "        [ 1.3335e-02,  3.8377e-02, -4.0418e-02,  ..., -1.1633e-02,\n",
+      "         -1.8374e-02, -1.1622e-02],\n",
+      "        [-5.3741e-02,  7.6245e-03,  5.9155e-02,  ..., -1.7577e-02,\n",
+      "         -4.4678e-02, -1.1400e-02],\n",
+      "        ...,\n",
+      "        [-3.7519e-02,  7.7807e-03,  2.1634e-02,  ..., -1.2836e-02,\n",
+      "         -3.3665e-02, -3.3249e-05],\n",
+      "        [ 8.0464e-03, -1.4498e-03,  1.2061e-02,  ...,  2.6527e-02,\n",
+      "          2.2661e-03,  1.5226e-02],\n",
+      "        [ 1.3127e-02,  3.4243e-02, -7.6076e-03,  ..., -1.7030e-02,\n",
+      "          3.7331e-02, -2.0033e-02]], device='cuda:0'), 'base_model.model.encoder.block.5.layer.0.SelfAttention.q.lora_B.weight': tensor([[-0.0211,  0.0119,  0.0196,  ...,  0.0310,  0.0261,  0.0212],\n",
+      "        [ 0.0150,  0.0073, -0.0072,  ..., -0.0296, -0.0019, -0.0055],\n",
+      "        [ 0.0040, -0.0227, -0.0023,  ..., -0.0025, -0.0282, -0.0286],\n",
+      "        ...,\n",
+      "        [-0.0339,  0.0068,  0.0066,  ...,  0.0262,  0.0112,  0.0062],\n",
+      "        [ 0.0111,  0.0063, -0.0170,  ..., -0.0230, -0.0087,  0.0053],\n",
+      "        [-0.0182,  0.0041,  0.0143,  ...,  0.0118,  0.0134, -0.0034]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.5.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0077, -0.0111,  0.0294,  ..., -0.0395,  0.0136,  0.0262],\n",
+      "        [-0.0040, -0.0357, -0.0582,  ..., -0.0002, -0.0220, -0.0344],\n",
+      "        [ 0.0358,  0.0259,  0.0419,  ...,  0.0121,  0.0241, -0.0243],\n",
+      "        ...,\n",
+      "        [ 0.0032, -0.0115, -0.0457,  ...,  0.0290,  0.0213, -0.0060],\n",
+      "        [ 0.0011,  0.0236,  0.0375,  ...,  0.0009,  0.0075,  0.0216],\n",
+      "        [-0.0085,  0.0274,  0.0160,  ...,  0.0069, -0.0044,  0.0034]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.5.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0093,  0.0031,  0.0132,  ..., -0.0114,  0.0090,  0.0082],\n",
+      "        [ 0.0036,  0.0175, -0.0117,  ...,  0.0118, -0.0122, -0.0122],\n",
+      "        [-0.0102,  0.0050,  0.0037,  ...,  0.0009, -0.0026, -0.0070],\n",
+      "        ...,\n",
+      "        [ 0.0107, -0.0064,  0.0131,  ..., -0.0112,  0.0124,  0.0116],\n",
+      "        [-0.0078,  0.0084,  0.0060,  ...,  0.0016,  0.0031, -0.0009],\n",
+      "        [ 0.0077,  0.0074, -0.0095,  ...,  0.0063, -0.0104, -0.0140]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.6.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0052, -0.0190, -0.0309,  ...,  0.0368,  0.0119,  0.0045],\n",
+      "        [ 0.0246,  0.0021, -0.0215,  ...,  0.0109, -0.0025, -0.0697],\n",
+      "        [-0.0161,  0.0301, -0.0043,  ..., -0.0135, -0.0248, -0.0280],\n",
+      "        ...,\n",
+      "        [ 0.0573, -0.0103,  0.0023,  ...,  0.0115,  0.0555, -0.0128],\n",
+      "        [-0.0021,  0.0424, -0.0038,  ...,  0.0515,  0.0286,  0.0078],\n",
+      "        [-0.0103, -0.0092,  0.0028,  ...,  0.0229, -0.0282, -0.0122]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.6.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0185,  0.0261, -0.0120,  ...,  0.0008,  0.0034, -0.0148],\n",
+      "        [-0.0174,  0.0061, -0.0011,  ..., -0.0079,  0.0253,  0.0086],\n",
+      "        [ 0.0133, -0.0025, -0.0073,  ..., -0.0178,  0.0182,  0.0027],\n",
+      "        ...,\n",
+      "        [ 0.0103, -0.0069, -0.0184,  ..., -0.0013,  0.0038, -0.0002],\n",
+      "        [ 0.0133,  0.0024,  0.0006,  ..., -0.0015, -0.0018,  0.0221],\n",
+      "        [ 0.0162, -0.0035, -0.0131,  ..., -0.0185, -0.0240,  0.0108]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.6.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0123, -0.0133, -0.0479,  ...,  0.0003,  0.0049, -0.0105],\n",
+      "        [-0.0214,  0.0090,  0.0156,  ..., -0.0335, -0.0089, -0.0161],\n",
+      "        [ 0.0283, -0.0100,  0.0151,  ..., -0.0126, -0.0030, -0.0421],\n",
+      "        ...,\n",
+      "        [ 0.0319,  0.0161, -0.0138,  ...,  0.0119,  0.0346,  0.0224],\n",
+      "        [ 0.0075,  0.0111,  0.0164,  ...,  0.0003,  0.0058,  0.0415],\n",
+      "        [ 0.0122,  0.0280, -0.0096,  ...,  0.0193,  0.0315,  0.0293]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.6.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0163,  0.0121,  0.0083,  ..., -0.0032, -0.0093,  0.0055],\n",
+      "        [ 0.0157, -0.0087, -0.0007,  ...,  0.0037,  0.0028,  0.0137],\n",
+      "        [ 0.0043, -0.0024, -0.0150,  ...,  0.0119,  0.0082,  0.0046],\n",
+      "        ...,\n",
+      "        [ 0.0077, -0.0089, -0.0054,  ...,  0.0100,  0.0072,  0.0122],\n",
+      "        [ 0.0042, -0.0140,  0.0039,  ...,  0.0058, -0.0015, -0.0027],\n",
+      "        [ 0.0118, -0.0002, -0.0154,  ..., -0.0005, -0.0107, -0.0160]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.7.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0249, -0.0245,  0.0263,  ...,  0.0231,  0.0002, -0.0161],\n",
+      "        [-0.0284,  0.0376, -0.0152,  ...,  0.0037, -0.0273, -0.0441],\n",
+      "        [ 0.0552, -0.0126,  0.0192,  ...,  0.0180,  0.0117, -0.0234],\n",
+      "        ...,\n",
+      "        [ 0.0255, -0.0163,  0.0113,  ..., -0.0212, -0.0115,  0.0077],\n",
+      "        [-0.0227, -0.0103, -0.0073,  ...,  0.0213, -0.0048, -0.0216],\n",
+      "        [ 0.0301, -0.0359, -0.0198,  ...,  0.0291,  0.0122,  0.0169]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.7.layer.0.SelfAttention.q.lora_B.weight': tensor([[-3.6794e-03, -1.1577e-02, -6.4594e-03,  ..., -5.8575e-03,\n",
+      "          6.4691e-03,  6.1957e-03],\n",
+      "        [-2.4289e-02,  4.5173e-03, -2.1506e-02,  ..., -2.8095e-02,\n",
+      "          4.1102e-02, -1.3792e-02],\n",
+      "        [ 7.2917e-03, -2.7127e-02,  6.7668e-05,  ...,  7.6069e-03,\n",
+      "         -3.3562e-03,  7.5219e-03],\n",
+      "        ...,\n",
+      "        [-4.4854e-03, -2.3905e-03, -6.4449e-03,  ..., -7.4162e-03,\n",
+      "          9.0684e-03,  2.0679e-04],\n",
+      "        [-4.0559e-03,  1.3454e-02, -5.7965e-03,  ...,  4.8874e-04,\n",
+      "          1.2333e-02, -8.3721e-03],\n",
+      "        [-3.8800e-03, -1.0870e-02, -1.6102e-02,  ..., -2.2499e-02,\n",
+      "          1.5781e-02,  1.3765e-02]], device='cuda:0'), 'base_model.model.encoder.block.7.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0115,  0.0034,  0.0256,  ..., -0.0181, -0.0146, -0.0219],\n",
+      "        [ 0.0442, -0.0186, -0.0151,  ..., -0.0228,  0.0146,  0.0321],\n",
+      "        [-0.0081, -0.0252,  0.0188,  ..., -0.0297,  0.0230,  0.0064],\n",
+      "        ...,\n",
+      "        [ 0.0098,  0.0081,  0.0102,  ..., -0.0116, -0.0435,  0.0061],\n",
+      "        [-0.0069,  0.0040, -0.0126,  ..., -0.0046, -0.0336,  0.0175],\n",
+      "        [ 0.0045, -0.0129,  0.0252,  ..., -0.0328, -0.0095, -0.0250]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.7.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0095, -0.0049,  0.0090,  ..., -0.0012, -0.0114,  0.0152],\n",
+      "        [-0.0096, -0.0163,  0.0093,  ..., -0.0169, -0.0110,  0.0143],\n",
+      "        [ 0.0172,  0.0033, -0.0209,  ...,  0.0011,  0.0101, -0.0149],\n",
+      "        ...,\n",
+      "        [-0.0223, -0.0201,  0.0178,  ...,  0.0135, -0.0135, -0.0085],\n",
+      "        [ 0.0029, -0.0050, -0.0139,  ..., -0.0083,  0.0040, -0.0070],\n",
+      "        [-0.0176, -0.0048,  0.0106,  ...,  0.0100, -0.0008,  0.0078]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.8.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0547, -0.0251, -0.0336,  ...,  0.0047, -0.0150, -0.0334],\n",
+      "        [ 0.0297,  0.0428, -0.0070,  ...,  0.0102,  0.0042, -0.0159],\n",
+      "        [-0.0193,  0.0235,  0.0417,  ..., -0.0160, -0.0147,  0.0220],\n",
+      "        ...,\n",
+      "        [ 0.0081, -0.0281,  0.0176,  ..., -0.0033,  0.0308,  0.0319],\n",
+      "        [ 0.0157,  0.0331, -0.0299,  ...,  0.0307,  0.0073,  0.0037],\n",
+      "        [-0.0328, -0.0115, -0.0006,  ..., -0.0026, -0.0101,  0.0336]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.8.layer.0.SelfAttention.q.lora_B.weight': tensor([[-0.0221,  0.0148,  0.0164,  ..., -0.0353, -0.0054, -0.0331],\n",
+      "        [ 0.0102,  0.0060, -0.0084,  ...,  0.0259,  0.0239,  0.0233],\n",
+      "        [-0.0187, -0.0233, -0.0057,  ..., -0.0188,  0.0055,  0.0035],\n",
+      "        ...,\n",
+      "        [ 0.0039, -0.0079, -0.0280,  ...,  0.0096,  0.0152,  0.0069],\n",
+      "        [-0.0040,  0.0051,  0.0182,  ..., -0.0046, -0.0015,  0.0005],\n",
+      "        [-0.0031,  0.0172, -0.0042,  ..., -0.0034,  0.0112,  0.0120]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.8.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0472,  0.0169, -0.0520,  ..., -0.0187,  0.0038,  0.0221],\n",
+      "        [-0.0078, -0.0035,  0.0112,  ..., -0.0010,  0.0342,  0.0090],\n",
+      "        [-0.0125,  0.0102,  0.0478,  ..., -0.0135, -0.0045,  0.0193],\n",
+      "        ...,\n",
+      "        [-0.0448, -0.0266,  0.0023,  ...,  0.0034,  0.0374,  0.0349],\n",
+      "        [ 0.0378,  0.0289, -0.0119,  ..., -0.0113,  0.0009, -0.0055],\n",
+      "        [ 0.0265, -0.0061, -0.0120,  ...,  0.0267, -0.0213, -0.0053]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.8.layer.0.SelfAttention.v.lora_B.weight': tensor([[-1.1674e-02, -1.4416e-02,  1.6741e-02,  ..., -2.0069e-02,\n",
+      "          1.3848e-02, -9.4953e-05],\n",
+      "        [ 2.8010e-03, -6.6800e-03, -1.4637e-03,  ..., -1.3755e-02,\n",
+      "          1.9308e-02,  1.2938e-02],\n",
+      "        [ 1.5742e-02, -7.6401e-03, -1.9585e-02,  ..., -7.9043e-03,\n",
+      "          9.8119e-03, -3.4114e-03],\n",
+      "        ...,\n",
+      "        [ 8.7199e-03, -1.3911e-02,  1.9556e-02,  ..., -1.1287e-02,\n",
+      "          1.3224e-02, -1.2388e-02],\n",
+      "        [-1.4302e-02,  2.5201e-03,  1.5830e-02,  ...,  1.1704e-02,\n",
+      "         -9.5458e-03, -2.3148e-03],\n",
+      "        [-1.2765e-02, -1.4806e-02,  3.0562e-02,  ..., -1.1230e-03,\n",
+      "         -2.1326e-03,  1.5665e-02]], device='cuda:0'), 'base_model.model.encoder.block.9.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0424,  0.0046, -0.0188,  ..., -0.0201,  0.0062, -0.0091],\n",
+      "        [ 0.0254,  0.0364,  0.0047,  ..., -0.0357, -0.0316, -0.0028],\n",
+      "        [-0.0442,  0.0029, -0.0119,  ..., -0.0363, -0.0380, -0.0143],\n",
+      "        ...,\n",
+      "        [ 0.0478,  0.0122,  0.0339,  ...,  0.0345, -0.0044, -0.0131],\n",
+      "        [ 0.0319,  0.0082, -0.0122,  ...,  0.0015, -0.0245, -0.0195],\n",
+      "        [ 0.0207,  0.0255,  0.0040,  ..., -0.0058, -0.0114, -0.0195]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.9.layer.0.SelfAttention.q.lora_B.weight': tensor([[-2.2374e-03,  5.3732e-03,  1.9213e-02,  ...,  1.1581e-02,\n",
+      "         -3.2477e-03, -4.9648e-04],\n",
+      "        [ 7.1751e-03, -8.7029e-03,  3.1709e-02,  ..., -2.3060e-02,\n",
+      "         -1.1236e-02,  1.3812e-02],\n",
+      "        [ 4.9562e-04, -4.9533e-03, -4.9173e-03,  ...,  2.0552e-02,\n",
+      "         -1.9678e-03, -3.8285e-05],\n",
+      "        ...,\n",
+      "        [ 1.4587e-02, -2.4020e-02, -1.1420e-02,  ...,  1.3391e-03,\n",
+      "         -4.4101e-03,  3.2557e-03],\n",
+      "        [-2.1797e-02,  3.1295e-02,  2.9931e-02,  ...,  1.0324e-02,\n",
+      "          2.1327e-02, -1.1652e-02],\n",
+      "        [-2.7132e-02,  5.5430e-03, -4.0748e-02,  ...,  7.3987e-03,\n",
+      "          1.2751e-02, -1.0383e-02]], device='cuda:0'), 'base_model.model.encoder.block.9.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0212,  0.0094, -0.0345,  ..., -0.0055,  0.0172,  0.0463],\n",
+      "        [-0.0218,  0.0271,  0.0145,  ..., -0.0182, -0.0750,  0.0068],\n",
+      "        [-0.0287, -0.0087,  0.0458,  ...,  0.0255,  0.0004, -0.0237],\n",
+      "        ...,\n",
+      "        [-0.0278,  0.0398,  0.0445,  ...,  0.0031, -0.0069, -0.0153],\n",
+      "        [ 0.0143, -0.0057,  0.0280,  ...,  0.0183, -0.0014,  0.0398],\n",
+      "        [-0.0347, -0.0115,  0.0330,  ...,  0.0281, -0.0414, -0.0309]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.9.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0175, -0.0150, -0.0187,  ..., -0.0194, -0.0235, -0.0187],\n",
+      "        [-0.0096,  0.0072,  0.0135,  ...,  0.0082,  0.0054,  0.0180],\n",
+      "        [-0.0190,  0.0146,  0.0120,  ...,  0.0130,  0.0177,  0.0227],\n",
+      "        ...,\n",
+      "        [-0.0145, -0.0077,  0.0004,  ...,  0.0100, -0.0152,  0.0171],\n",
+      "        [ 0.0154, -0.0045, -0.0016,  ..., -0.0170,  0.0122, -0.0131],\n",
+      "        [-0.0150, -0.0046, -0.0022,  ...,  0.0112, -0.0164,  0.0145]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.10.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0063, -0.0355,  0.0021,  ...,  0.0098,  0.0202, -0.0121],\n",
+      "        [-0.0238,  0.0037,  0.0345,  ..., -0.0088, -0.0315, -0.0167],\n",
+      "        [ 0.0099, -0.0115, -0.0161,  ..., -0.0287,  0.0219,  0.0324],\n",
+      "        ...,\n",
+      "        [ 0.0052, -0.0403, -0.0258,  ..., -0.0323, -0.0180,  0.0049],\n",
+      "        [ 0.0460, -0.0125,  0.0313,  ...,  0.0127,  0.0384, -0.0161],\n",
+      "        [ 0.0363, -0.0080, -0.0062,  ...,  0.0132,  0.0386, -0.0175]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.10.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0137, -0.0086, -0.0250,  ..., -0.0072, -0.0013,  0.0127],\n",
+      "        [ 0.0238,  0.0095, -0.0219,  ...,  0.0205,  0.0050,  0.0102],\n",
+      "        [ 0.0054,  0.0194, -0.0351,  ..., -0.0171,  0.0097,  0.0261],\n",
+      "        ...,\n",
+      "        [-0.0058, -0.0221,  0.0088,  ..., -0.0028,  0.0017, -0.0202],\n",
+      "        [ 0.0044, -0.0364,  0.0107,  ..., -0.0187, -0.0172, -0.0214],\n",
+      "        [-0.0165, -0.0024,  0.0094,  ..., -0.0130, -0.0179, -0.0093]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.10.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0242, -0.0273,  0.0012,  ..., -0.0284, -0.0053, -0.0016],\n",
+      "        [ 0.0436, -0.0018, -0.0349,  ...,  0.0164,  0.0422,  0.0017],\n",
+      "        [ 0.0485,  0.0392,  0.0135,  ...,  0.0272, -0.0033, -0.0040],\n",
+      "        ...,\n",
+      "        [ 0.0098,  0.0046, -0.0287,  ..., -0.0284,  0.0120,  0.0047],\n",
+      "        [ 0.0139, -0.0017, -0.0226,  ..., -0.0269, -0.0447,  0.0190],\n",
+      "        [ 0.0453, -0.0130, -0.0016,  ...,  0.0023,  0.0084,  0.0260]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.10.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0052, -0.0034,  0.0169,  ..., -0.0010,  0.0101,  0.0065],\n",
+      "        [-0.0235,  0.0224, -0.0020,  ...,  0.0105, -0.0222, -0.0214],\n",
+      "        [-0.0252,  0.0259, -0.0122,  ...,  0.0191, -0.0154, -0.0236],\n",
+      "        ...,\n",
+      "        [-0.0093,  0.0063,  0.0104,  ...,  0.0207, -0.0164, -0.0046],\n",
+      "        [-0.0115,  0.0047, -0.0123,  ...,  0.0146, -0.0010, -0.0085],\n",
+      "        [-0.0081,  0.0099,  0.0202,  ...,  0.0214, -0.0252,  0.0029]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.11.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0178, -0.0056,  0.0114,  ...,  0.0069,  0.0038,  0.0309],\n",
+      "        [ 0.0363,  0.0298,  0.0199,  ..., -0.0196, -0.0228,  0.0145],\n",
+      "        [-0.0566,  0.0105, -0.0255,  ..., -0.0212,  0.0047,  0.0039],\n",
+      "        ...,\n",
+      "        [ 0.0267,  0.0304, -0.0092,  ...,  0.0117, -0.0274, -0.0169],\n",
+      "        [ 0.0403,  0.0225, -0.0204,  ..., -0.0277,  0.0124, -0.0060],\n",
+      "        [ 0.0238, -0.0275,  0.0223,  ...,  0.0041, -0.0117,  0.0226]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.11.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 1.5251e-02, -1.3405e-02, -1.0952e-02,  ..., -3.2804e-02,\n",
+      "          8.0364e-03,  1.1897e-02],\n",
+      "        [-1.2090e-02,  1.9948e-02,  2.2969e-03,  ...,  1.9390e-02,\n",
+      "          5.5107e-03,  8.3068e-03],\n",
+      "        [-9.2829e-03,  1.3713e-02,  2.2380e-02,  ..., -3.9204e-03,\n",
+      "         -1.3160e-02, -3.6896e-03],\n",
+      "        ...,\n",
+      "        [-8.0247e-03,  6.4117e-03,  1.9469e-02,  ...,  1.0616e-02,\n",
+      "          2.5570e-04,  2.6831e-03],\n",
+      "        [-5.6667e-03, -5.0435e-03, -1.5369e-02,  ..., -1.1559e-02,\n",
+      "          1.9381e-02, -7.5156e-05],\n",
+      "        [ 5.1334e-03,  9.1880e-03, -2.8864e-02,  ..., -2.4764e-02,\n",
+      "          1.5606e-02, -2.5908e-02]], device='cuda:0'), 'base_model.model.encoder.block.11.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0133, -0.0197, -0.0365,  ..., -0.0401, -0.0523,  0.0024],\n",
+      "        [-0.0132,  0.0168,  0.0084,  ...,  0.0424, -0.0226,  0.0277],\n",
+      "        [ 0.0298,  0.0316, -0.0208,  ...,  0.0038,  0.0429,  0.0068],\n",
+      "        ...,\n",
+      "        [ 0.0310,  0.0084,  0.0016,  ...,  0.0036,  0.0146, -0.0063],\n",
+      "        [-0.0516, -0.0046, -0.0126,  ...,  0.0213, -0.0387, -0.0145],\n",
+      "        [-0.0594, -0.0328, -0.0431,  ..., -0.0225, -0.0079, -0.0229]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.11.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0250, -0.0376,  0.0345,  ...,  0.0266, -0.0233, -0.0127],\n",
+      "        [-0.0047,  0.0099, -0.0125,  ..., -0.0094,  0.0040,  0.0016],\n",
+      "        [-0.0066,  0.0164, -0.0167,  ..., -0.0113,  0.0158,  0.0106],\n",
+      "        ...,\n",
+      "        [-0.0111,  0.0109, -0.0093,  ..., -0.0007,  0.0101, -0.0061],\n",
+      "        [-0.0109,  0.0009, -0.0044,  ..., -0.0191, -0.0032,  0.0205],\n",
+      "        [-0.0100,  0.0034, -0.0044,  ..., -0.0122,  0.0104,  0.0291]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.12.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0140,  0.0012, -0.0323,  ...,  0.0490, -0.0234, -0.0137],\n",
+      "        [-0.0366, -0.0191,  0.0091,  ...,  0.0381,  0.0512, -0.0114],\n",
+      "        [ 0.0090, -0.0214, -0.0278,  ...,  0.0384,  0.0307, -0.0009],\n",
+      "        ...,\n",
+      "        [-0.0556,  0.0305,  0.0202,  ...,  0.0185, -0.0310, -0.0131],\n",
+      "        [-0.0287,  0.0159,  0.0026,  ...,  0.0202,  0.0368,  0.0184],\n",
+      "        [-0.0384,  0.0130, -0.0241,  ...,  0.0191,  0.0524,  0.0031]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.12.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0059,  0.0123,  0.0067,  ...,  0.0095,  0.0090,  0.0237],\n",
+      "        [ 0.0024,  0.0133,  0.0313,  ...,  0.0079,  0.0124,  0.0065],\n",
+      "        [ 0.0086, -0.0328, -0.0109,  ..., -0.0251, -0.0013, -0.0018],\n",
+      "        ...,\n",
+      "        [-0.0003, -0.0192,  0.0292,  ..., -0.0260,  0.0006, -0.0082],\n",
+      "        [ 0.0058, -0.0187, -0.0136,  ..., -0.0056, -0.0078, -0.0284],\n",
+      "        [ 0.0213, -0.0099, -0.0119,  ...,  0.0173,  0.0065, -0.0007]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.12.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0155,  0.0088,  0.0113,  ...,  0.0075,  0.0049,  0.0154],\n",
+      "        [-0.0268,  0.0203,  0.0017,  ..., -0.0269,  0.0272,  0.0151],\n",
+      "        [-0.0063,  0.0055, -0.0013,  ..., -0.0193,  0.0099, -0.0143],\n",
+      "        ...,\n",
+      "        [ 0.0315,  0.0117, -0.0101,  ...,  0.0034,  0.0140,  0.0136],\n",
+      "        [ 0.0042, -0.0548, -0.0144,  ...,  0.0041,  0.0296,  0.0246],\n",
+      "        [ 0.0081,  0.0255,  0.0194,  ...,  0.0020,  0.0124,  0.0246]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.12.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0034,  0.0007, -0.0139,  ..., -0.0021,  0.0076,  0.0024],\n",
+      "        [-0.0063, -0.0191, -0.0012,  ...,  0.0197,  0.0059,  0.0199],\n",
+      "        [ 0.0098,  0.0127, -0.0077,  ..., -0.0152, -0.0014, -0.0179],\n",
+      "        ...,\n",
+      "        [-0.0012, -0.0055,  0.0102,  ...,  0.0071,  0.0157,  0.0007],\n",
+      "        [-0.0223, -0.0310, -0.0248,  ...,  0.0279,  0.0295,  0.0272],\n",
+      "        [-0.0144, -0.0010, -0.0043,  ...,  0.0101,  0.0027,  0.0096]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.13.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0528,  0.0265, -0.0006,  ...,  0.0108,  0.0268, -0.0358],\n",
+      "        [-0.0148, -0.0179,  0.0211,  ..., -0.0519, -0.0323, -0.0162],\n",
+      "        [-0.0082, -0.0202, -0.0363,  ..., -0.0067,  0.0422,  0.0203],\n",
+      "        ...,\n",
+      "        [ 0.0119,  0.0068,  0.0021,  ..., -0.0284, -0.0126, -0.0077],\n",
+      "        [-0.0252,  0.0406, -0.0096,  ...,  0.0380, -0.0298, -0.0124],\n",
+      "        [ 0.0174, -0.0255,  0.0100,  ..., -0.0062, -0.0038, -0.0483]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.13.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0046,  0.0007,  0.0040,  ...,  0.0127, -0.0086,  0.0045],\n",
+      "        [-0.0034,  0.0012, -0.0269,  ...,  0.0077,  0.0077,  0.0079],\n",
+      "        [-0.0255,  0.0091,  0.0121,  ...,  0.0027, -0.0190,  0.0250],\n",
+      "        ...,\n",
+      "        [-0.0013, -0.0045,  0.0094,  ..., -0.0040, -0.0030,  0.0165],\n",
+      "        [-0.0094, -0.0013,  0.0098,  ...,  0.0043, -0.0060,  0.0057],\n",
+      "        [-0.0247,  0.0240,  0.0263,  ...,  0.0149, -0.0302,  0.0101]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.13.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0235,  0.0104,  0.0366,  ...,  0.0335,  0.0118, -0.0104],\n",
+      "        [-0.0169,  0.0107, -0.0293,  ..., -0.0196,  0.0248,  0.0136],\n",
+      "        [ 0.0046,  0.0114,  0.0013,  ..., -0.0061,  0.0208,  0.0137],\n",
+      "        ...,\n",
+      "        [-0.0008,  0.0332,  0.0251,  ..., -0.0024,  0.0390, -0.0130],\n",
+      "        [ 0.0226, -0.0223, -0.0490,  ..., -0.0123, -0.0358,  0.0141],\n",
+      "        [-0.0114, -0.0098, -0.0323,  ...,  0.0350, -0.0451,  0.0060]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.13.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 1.3388e-03, -7.8071e-03, -4.4231e-03,  ...,  6.8468e-03,\n",
+      "         -8.6865e-03, -6.6189e-03],\n",
+      "        [-8.4138e-03,  3.9207e-03, -1.7379e-02,  ..., -1.4644e-02,\n",
+      "          2.0210e-02,  1.8017e-02],\n",
+      "        [-8.6932e-03,  3.4929e-03, -8.5747e-03,  ..., -7.9256e-05,\n",
+      "         -7.3968e-03, -1.2425e-02],\n",
+      "        ...,\n",
+      "        [ 2.5251e-02, -1.5215e-02,  1.2540e-02,  ...,  5.3171e-04,\n",
+      "          5.1631e-03, -6.5597e-04],\n",
+      "        [-4.4899e-03, -1.2352e-02, -3.8436e-03,  ...,  4.8080e-03,\n",
+      "         -9.7884e-04, -6.4385e-03],\n",
+      "        [ 1.2952e-02, -6.3606e-03,  6.3772e-03,  ..., -9.4017e-03,\n",
+      "          5.4179e-03,  6.1516e-03]], device='cuda:0'), 'base_model.model.encoder.block.14.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0031, -0.0145, -0.0498,  ...,  0.0126, -0.0046,  0.0666],\n",
+      "        [-0.0352,  0.0089,  0.0178,  ..., -0.0369, -0.0284, -0.0033],\n",
+      "        [-0.0229,  0.0080, -0.0020,  ...,  0.0681,  0.0048,  0.0553],\n",
+      "        ...,\n",
+      "        [ 0.0118,  0.0116, -0.0289,  ...,  0.0028, -0.0169,  0.0415],\n",
+      "        [ 0.0137, -0.0024,  0.0268,  ...,  0.0435,  0.0123, -0.0126],\n",
+      "        [ 0.0422,  0.0337,  0.0022,  ...,  0.0373,  0.0318,  0.0344]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.14.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0116, -0.0045,  0.0094,  ...,  0.0165,  0.0021,  0.0011],\n",
+      "        [-0.0034,  0.0033, -0.0199,  ...,  0.0003,  0.0146,  0.0141],\n",
+      "        [-0.0055,  0.0080, -0.0076,  ..., -0.0197, -0.0134, -0.0145],\n",
+      "        ...,\n",
+      "        [ 0.0130, -0.0213,  0.0119,  ...,  0.0132,  0.0252,  0.0126],\n",
+      "        [ 0.0011, -0.0145,  0.0209,  ..., -0.0021,  0.0018, -0.0081],\n",
+      "        [-0.0271, -0.0154,  0.0062,  ...,  0.0095,  0.0025, -0.0112]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.14.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0097, -0.0242,  0.0057,  ...,  0.0214, -0.0358, -0.0250],\n",
+      "        [-0.0355, -0.0014, -0.0570,  ...,  0.0239,  0.0224, -0.0252],\n",
+      "        [-0.0124,  0.0291,  0.0428,  ..., -0.0319,  0.0109,  0.0128],\n",
+      "        ...,\n",
+      "        [ 0.0115, -0.0210, -0.0035,  ...,  0.0227, -0.0120,  0.0143],\n",
+      "        [-0.0309,  0.0038,  0.0248,  ...,  0.0052, -0.0469,  0.0364],\n",
+      "        [-0.0103,  0.0015, -0.0015,  ..., -0.0410,  0.0415, -0.0185]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.14.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0389, -0.0082,  0.0113,  ..., -0.0270,  0.0057,  0.0258],\n",
+      "        [ 0.0010,  0.0043, -0.0075,  ..., -0.0001, -0.0128, -0.0028],\n",
+      "        [ 0.0128,  0.0135, -0.0123,  ...,  0.0104, -0.0072, -0.0109],\n",
+      "        ...,\n",
+      "        [-0.0414, -0.0318,  0.0169,  ..., -0.0431,  0.0356,  0.0451],\n",
+      "        [ 0.0091, -0.0132,  0.0102,  ..., -0.0135,  0.0111, -0.0035],\n",
+      "        [ 0.0005,  0.0068, -0.0088,  ..., -0.0002,  0.0083, -0.0006]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.15.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0069,  0.0003, -0.0256,  ...,  0.0258, -0.0335,  0.0484],\n",
+      "        [-0.0216, -0.0322,  0.0204,  ...,  0.0170, -0.0222, -0.0108],\n",
+      "        [ 0.0008,  0.0105, -0.0380,  ..., -0.0099, -0.0394,  0.0262],\n",
+      "        ...,\n",
+      "        [-0.0015, -0.0016, -0.0252,  ..., -0.0252, -0.0368, -0.0343],\n",
+      "        [-0.0253,  0.0274,  0.0133,  ...,  0.0494, -0.0482,  0.0473],\n",
+      "        [-0.0317,  0.0028,  0.0004,  ...,  0.0471,  0.0196,  0.0111]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.15.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0266, -0.0135,  0.0290,  ..., -0.0318, -0.0019,  0.0112],\n",
+      "        [-0.0016,  0.0150, -0.0023,  ...,  0.0111, -0.0071, -0.0280],\n",
+      "        [-0.0286,  0.0020, -0.0203,  ...,  0.0246, -0.0086, -0.0261],\n",
+      "        ...,\n",
+      "        [-0.0133, -0.0072, -0.0033,  ...,  0.0133, -0.0128, -0.0120],\n",
+      "        [-0.0055, -0.0159, -0.0040,  ...,  0.0039, -0.0175, -0.0134],\n",
+      "        [-0.0024,  0.0035,  0.0067,  ...,  0.0042, -0.0094, -0.0038]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.15.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 1.7631e-03,  3.6335e-02, -2.2382e-02,  ..., -1.4592e-02,\n",
+      "          5.8360e-02, -1.8392e-02],\n",
+      "        [ 1.4876e-02,  2.0355e-02, -2.3993e-02,  ...,  1.8249e-02,\n",
+      "         -2.0012e-03, -2.4304e-02],\n",
+      "        [-1.7222e-02,  8.3348e-04, -2.3049e-02,  ..., -2.6813e-02,\n",
+      "          3.8505e-02, -1.4321e-02],\n",
+      "        ...,\n",
+      "        [-1.3301e-02, -1.7976e-05, -1.0527e-02,  ...,  4.7702e-02,\n",
+      "          2.2238e-03, -2.4194e-02],\n",
+      "        [-2.4366e-02, -3.5398e-02,  2.0618e-02,  ...,  3.4942e-03,\n",
+      "         -4.9058e-02,  8.0171e-03],\n",
+      "        [-3.0048e-03,  3.7638e-02, -3.8622e-03,  ..., -5.0610e-03,\n",
+      "          5.2482e-02, -1.0379e-02]], device='cuda:0'), 'base_model.model.encoder.block.15.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 2.5725e-03,  3.5926e-03,  1.9242e-03,  ..., -1.1365e-02,\n",
+      "         -2.6424e-04, -4.3521e-03],\n",
+      "        [-2.1243e-03, -1.0515e-02,  4.6888e-03,  ..., -3.1430e-03,\n",
+      "         -6.3100e-03, -6.4766e-03],\n",
+      "        [-1.0521e-02, -1.0788e-02, -6.9606e-03,  ...,  1.1522e-02,\n",
+      "          3.2694e-03, -1.3453e-02],\n",
+      "        ...,\n",
+      "        [-3.5505e-05, -2.4789e-02,  5.3724e-03,  ...,  1.2047e-02,\n",
+      "         -5.5704e-03, -1.3944e-02],\n",
+      "        [-1.4204e-02, -2.3716e-02, -2.5142e-03,  ...,  1.1288e-02,\n",
+      "          1.0336e-02, -1.5918e-02],\n",
+      "        [-9.6600e-03, -4.7291e-04, -1.8597e-02,  ...,  1.7493e-02,\n",
+      "          1.8787e-02, -1.6483e-02]], device='cuda:0'), 'base_model.model.encoder.block.16.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0099, -0.0129,  0.0517,  ..., -0.0158,  0.0089,  0.0313],\n",
+      "        [ 0.0373, -0.0277,  0.0096,  ..., -0.0250, -0.0250,  0.0282],\n",
+      "        [-0.0122, -0.0129, -0.0159,  ..., -0.0003,  0.0164,  0.0209],\n",
+      "        ...,\n",
+      "        [ 0.0360, -0.0043,  0.0127,  ...,  0.0088, -0.0333, -0.0175],\n",
+      "        [ 0.0071, -0.0533, -0.0338,  ...,  0.0417, -0.0052,  0.0208],\n",
+      "        [ 0.0142, -0.0381,  0.0059,  ..., -0.0328,  0.0163,  0.0182]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.16.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0117, -0.0307, -0.0021,  ...,  0.0123, -0.0253, -0.0016],\n",
+      "        [-0.0085,  0.0046,  0.0009,  ..., -0.0039,  0.0105, -0.0178],\n",
+      "        [ 0.0037,  0.0178,  0.0198,  ..., -0.0141,  0.0087,  0.0110],\n",
+      "        ...,\n",
+      "        [-0.0101,  0.0057,  0.0065,  ...,  0.0019,  0.0080, -0.0162],\n",
+      "        [ 0.0104,  0.0138,  0.0215,  ..., -0.0142,  0.0151,  0.0108],\n",
+      "        [ 0.0186,  0.0051,  0.0312,  ..., -0.0250,  0.0252,  0.0021]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.16.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0258, -0.0153, -0.0037,  ..., -0.0136, -0.0433,  0.0465],\n",
+      "        [-0.0139, -0.0489, -0.0567,  ...,  0.0056,  0.0010,  0.0062],\n",
+      "        [ 0.0086,  0.0541,  0.0437,  ..., -0.0202,  0.0503, -0.0420],\n",
+      "        ...,\n",
+      "        [ 0.0304,  0.0422, -0.0691,  ..., -0.0162,  0.0146,  0.0323],\n",
+      "        [-0.0310,  0.0305,  0.0252,  ...,  0.0340, -0.0040, -0.0321],\n",
+      "        [ 0.0063,  0.0173, -0.0269,  ...,  0.0041, -0.0490,  0.0162]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.16.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0372, -0.0004, -0.0395,  ..., -0.0252, -0.0212,  0.0400],\n",
+      "        [ 0.0557,  0.0030, -0.0577,  ..., -0.0191, -0.0292,  0.0516],\n",
+      "        [-0.0057,  0.0290,  0.0045,  ...,  0.0348, -0.0138, -0.0093],\n",
+      "        ...,\n",
+      "        [-0.0273, -0.0235,  0.0251,  ..., -0.0232,  0.0207, -0.0205],\n",
+      "        [ 0.0374,  0.0416, -0.0401,  ...,  0.0239, -0.0532,  0.0429],\n",
+      "        [-0.0201, -0.0159,  0.0268,  ..., -0.0145,  0.0342, -0.0210]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.17.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0124,  0.0274,  0.0223,  ..., -0.0135,  0.0262,  0.0426],\n",
+      "        [-0.0069,  0.0114,  0.0237,  ...,  0.0522,  0.0042,  0.0107],\n",
+      "        [ 0.0016, -0.0762, -0.0267,  ...,  0.0016,  0.0416, -0.0379],\n",
+      "        ...,\n",
+      "        [-0.0107, -0.0051,  0.0061,  ...,  0.0077, -0.0267, -0.0538],\n",
+      "        [-0.0166, -0.0188, -0.0361,  ..., -0.0321,  0.0282, -0.0272],\n",
+      "        [ 0.0195, -0.0308,  0.0041,  ..., -0.0343, -0.0432, -0.0337]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.17.layer.0.SelfAttention.q.lora_B.weight': tensor([[-1.4683e-02, -9.0315e-03,  1.3334e-02,  ...,  6.3660e-03,\n",
+      "          3.7701e-03,  1.8145e-02],\n",
+      "        [ 2.2848e-03, -4.4163e-03,  1.6257e-02,  ...,  8.8848e-03,\n",
+      "          1.0144e-02,  6.2522e-03],\n",
+      "        [ 1.5033e-02,  1.0294e-02, -1.7981e-02,  ..., -8.9641e-03,\n",
+      "         -2.1745e-02, -2.4771e-02],\n",
+      "        ...,\n",
+      "        [ 2.1272e-04, -7.2553e-03, -2.3124e-06,  ...,  3.4841e-03,\n",
+      "         -1.0348e-02, -2.2534e-02],\n",
+      "        [-2.2342e-03,  1.8137e-02, -6.5606e-03,  ...,  2.7981e-03,\n",
+      "         -8.4023e-03,  1.4239e-02],\n",
+      "        [-3.0212e-03,  3.7128e-03, -1.7892e-02,  ...,  1.1220e-02,\n",
+      "         -2.2024e-02, -8.1391e-03]], device='cuda:0'), 'base_model.model.encoder.block.17.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0019, -0.0137, -0.0012,  ..., -0.0406, -0.0212,  0.0466],\n",
+      "        [ 0.0227, -0.0154,  0.0074,  ...,  0.0155, -0.0104,  0.0346],\n",
+      "        [-0.0455,  0.0035,  0.0101,  ..., -0.0231, -0.0183, -0.0431],\n",
+      "        ...,\n",
+      "        [-0.0447,  0.0181,  0.0305,  ...,  0.0417, -0.0317, -0.0054],\n",
+      "        [ 0.0427, -0.0357, -0.0128,  ..., -0.0214,  0.0126,  0.0405],\n",
+      "        [-0.0563,  0.0184, -0.0125,  ...,  0.0406, -0.0077, -0.0211]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.17.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0090, -0.0032, -0.0055,  ..., -0.0054, -0.0003,  0.0061],\n",
+      "        [ 0.0058,  0.0085, -0.0087,  ..., -0.0043,  0.0099,  0.0082],\n",
+      "        [ 0.0087,  0.0208, -0.0196,  ..., -0.0181,  0.0206,  0.0036],\n",
+      "        ...,\n",
+      "        [-0.0395, -0.0220,  0.0149,  ...,  0.0223, -0.0172,  0.0448],\n",
+      "        [-0.0227, -0.0071,  0.0050,  ...,  0.0142, -0.0116,  0.0192],\n",
+      "        [-0.0102,  0.0058, -0.0205,  ..., -0.0027,  0.0044,  0.0175]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.18.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0241,  0.0187,  0.0110,  ...,  0.0043, -0.0370,  0.0384],\n",
+      "        [ 0.0038,  0.0020, -0.0249,  ..., -0.0397,  0.0138, -0.0315],\n",
+      "        [-0.0226, -0.0243, -0.0236,  ...,  0.0053, -0.0367, -0.0319],\n",
+      "        ...,\n",
+      "        [ 0.0032,  0.0287, -0.0235,  ..., -0.0466, -0.0053,  0.0106],\n",
+      "        [-0.0264, -0.0175, -0.0500,  ..., -0.0234,  0.0288, -0.0165],\n",
+      "        [-0.0156,  0.0258,  0.0038,  ..., -0.0575, -0.0257, -0.0088]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.18.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0184, -0.0091, -0.0198,  ..., -0.0144, -0.0019, -0.0017],\n",
+      "        [-0.0016,  0.0060, -0.0113,  ...,  0.0096, -0.0067,  0.0017],\n",
+      "        [-0.0023, -0.0006, -0.0505,  ..., -0.0005, -0.0123,  0.0116],\n",
+      "        ...,\n",
+      "        [-0.0247,  0.0048,  0.0250,  ...,  0.0066, -0.0144,  0.0116],\n",
+      "        [ 0.0060, -0.0079, -0.0068,  ..., -0.0094,  0.0171, -0.0095],\n",
+      "        [-0.0086, -0.0147,  0.0030,  ..., -0.0247,  0.0092, -0.0146]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.18.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0043,  0.0391, -0.0416,  ..., -0.0416,  0.0277,  0.0104],\n",
+      "        [-0.0030,  0.0658, -0.0103,  ..., -0.0245,  0.0268,  0.0276],\n",
+      "        [ 0.0047,  0.0221, -0.0067,  ...,  0.0026,  0.0197,  0.0375],\n",
+      "        ...,\n",
+      "        [-0.0245, -0.0708,  0.0604,  ...,  0.0216,  0.0002, -0.0149],\n",
+      "        [ 0.0146, -0.0349,  0.0189,  ..., -0.0136, -0.0235,  0.0123],\n",
+      "        [-0.0083, -0.0331,  0.0403,  ...,  0.0157, -0.0047, -0.0127]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.18.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0209, -0.0291, -0.0244,  ...,  0.0253,  0.0127, -0.0248],\n",
+      "        [ 0.0053,  0.0060,  0.0070,  ..., -0.0051, -0.0021, -0.0103],\n",
+      "        [ 0.0148,  0.0031,  0.0126,  ..., -0.0117, -0.0017, -0.0147],\n",
+      "        ...,\n",
+      "        [ 0.0200,  0.0158,  0.0231,  ..., -0.0149, -0.0121,  0.0300],\n",
+      "        [-0.0008, -0.0023, -0.0031,  ...,  0.0107,  0.0033,  0.0051],\n",
+      "        [ 0.0240,  0.0208,  0.0184,  ..., -0.0156, -0.0005,  0.0188]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.19.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0249, -0.0210, -0.0265,  ...,  0.0108, -0.0111, -0.0171],\n",
+      "        [-0.0351,  0.0088, -0.0190,  ...,  0.0171,  0.0221,  0.0042],\n",
+      "        [ 0.0030,  0.0014, -0.0108,  ...,  0.0105,  0.0679,  0.0146],\n",
+      "        ...,\n",
+      "        [ 0.0347, -0.0230,  0.0132,  ..., -0.0003, -0.0530,  0.0031],\n",
+      "        [ 0.0138,  0.0345,  0.0108,  ...,  0.0178,  0.0252,  0.0180],\n",
+      "        [ 0.0034,  0.0017, -0.0196,  ..., -0.0089, -0.0449, -0.0323]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.19.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0205, -0.0094, -0.0139,  ...,  0.0070, -0.0126,  0.0226],\n",
+      "        [-0.0010, -0.0070, -0.0134,  ...,  0.0123, -0.0060,  0.0030],\n",
+      "        [-0.0007, -0.0004, -0.0068,  ...,  0.0005,  0.0003,  0.0037],\n",
+      "        ...,\n",
+      "        [ 0.0378, -0.0430, -0.0296,  ...,  0.0429, -0.0451,  0.0347],\n",
+      "        [-0.0320,  0.0106,  0.0203,  ..., -0.0211,  0.0239, -0.0175],\n",
+      "        [ 0.0210, -0.0117, -0.0150,  ...,  0.0160, -0.0256,  0.0168]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.19.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 1.4258e-02, -2.0162e-02,  1.4808e-02,  ...,  6.3350e-02,\n",
+      "          3.1473e-02, -2.1102e-02],\n",
+      "        [ 3.0901e-03, -4.1341e-02,  2.0951e-02,  ...,  1.8360e-02,\n",
+      "         -6.7527e-03,  1.2380e-02],\n",
+      "        [ 2.9749e-02,  5.5730e-05, -6.0016e-03,  ...,  4.4938e-02,\n",
+      "          6.8638e-03, -1.5086e-02],\n",
+      "        ...,\n",
+      "        [-6.1253e-03,  1.3493e-03,  1.6801e-02,  ...,  5.0317e-03,\n",
+      "         -5.3136e-03,  1.7972e-02],\n",
+      "        [ 1.9262e-02, -1.0739e-03,  5.0379e-02,  ...,  6.5103e-02,\n",
+      "          2.6762e-02, -1.0854e-02],\n",
+      "        [ 2.7448e-02,  2.8939e-03,  4.7538e-02,  ...,  4.8242e-02,\n",
+      "          2.4452e-02, -2.7716e-02]], device='cuda:0'), 'base_model.model.encoder.block.19.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0130, -0.0170, -0.0146,  ..., -0.0144, -0.0169, -0.0175],\n",
+      "        [-0.0181, -0.0188, -0.0173,  ..., -0.0193, -0.0192, -0.0150],\n",
+      "        [ 0.0139,  0.0166,  0.0089,  ...,  0.0140,  0.0107,  0.0067],\n",
+      "        ...,\n",
+      "        [-0.0119, -0.0102, -0.0134,  ..., -0.0145, -0.0182, -0.0151],\n",
+      "        [ 0.0106,  0.0094,  0.0102,  ...,  0.0079,  0.0062,  0.0125],\n",
+      "        [-0.0259, -0.0304, -0.0267,  ..., -0.0262, -0.0288, -0.0286]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.20.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0136, -0.0010, -0.0112,  ...,  0.0021, -0.0050,  0.0332],\n",
+      "        [ 0.0083, -0.0270, -0.0023,  ...,  0.0004, -0.0469,  0.0221],\n",
+      "        [-0.0129, -0.0622,  0.0074,  ...,  0.0102, -0.0384,  0.0262],\n",
+      "        ...,\n",
+      "        [-0.0207, -0.0472,  0.0020,  ...,  0.0109, -0.0551,  0.0010],\n",
+      "        [-0.0171,  0.0238, -0.0192,  ...,  0.0388,  0.0184, -0.0068],\n",
+      "        [ 0.0377, -0.0111,  0.0170,  ..., -0.0148, -0.0103,  0.0415]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.20.layer.0.SelfAttention.q.lora_B.weight': tensor([[-0.0201, -0.0002, -0.0038,  ..., -0.0076,  0.0105, -0.0084],\n",
+      "        [-0.0349,  0.0177, -0.0228,  ...,  0.0044,  0.0148, -0.0253],\n",
+      "        [-0.0230,  0.0061,  0.0187,  ...,  0.0196, -0.0039,  0.0043],\n",
+      "        ...,\n",
+      "        [-0.0062,  0.0211, -0.0158,  ..., -0.0114,  0.0234, -0.0168],\n",
+      "        [-0.0158, -0.0012,  0.0284,  ...,  0.0266, -0.0246,  0.0207],\n",
+      "        [-0.0249,  0.0377, -0.0102,  ...,  0.0059,  0.0032, -0.0116]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.20.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0300,  0.0260,  0.0026,  ...,  0.0304,  0.0345,  0.0474],\n",
+      "        [-0.0201, -0.0045,  0.0204,  ..., -0.0125, -0.0221, -0.0285],\n",
+      "        [-0.0180, -0.0343,  0.0173,  ...,  0.0428,  0.0268,  0.0424],\n",
+      "        ...,\n",
+      "        [ 0.0010, -0.0441, -0.0271,  ...,  0.0089,  0.0099, -0.0180],\n",
+      "        [-0.0098, -0.0109,  0.0484,  ...,  0.0118,  0.0068, -0.0187],\n",
+      "        [-0.0206, -0.0263, -0.0238,  ...,  0.0390,  0.0060,  0.0163]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.20.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0201, -0.0216,  0.0119,  ...,  0.0183, -0.0170,  0.0124],\n",
+      "        [ 0.0130, -0.0100,  0.0180,  ...,  0.0223,  0.0051,  0.0060],\n",
+      "        [-0.0278,  0.0233, -0.0264,  ..., -0.0253,  0.0271, -0.0330],\n",
+      "        ...,\n",
+      "        [-0.0052, -0.0034,  0.0056,  ...,  0.0010,  0.0099,  0.0036],\n",
+      "        [-0.0003, -0.0029, -0.0112,  ..., -0.0089, -0.0206,  0.0131],\n",
+      "        [ 0.0036, -0.0245,  0.0003,  ..., -0.0018, -0.0305,  0.0315]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.21.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0078, -0.0170,  0.0393,  ..., -0.0284,  0.0189,  0.0169],\n",
+      "        [-0.0003,  0.0058, -0.0220,  ...,  0.0204, -0.0052, -0.0126],\n",
+      "        [ 0.0164,  0.0175, -0.0063,  ...,  0.0019,  0.0210, -0.0453],\n",
+      "        ...,\n",
+      "        [ 0.0046,  0.0172,  0.0201,  ..., -0.0191,  0.0161,  0.0512],\n",
+      "        [ 0.0214, -0.0041,  0.0220,  ..., -0.0172,  0.0319,  0.0360],\n",
+      "        [-0.0295,  0.0370, -0.0062,  ...,  0.0288,  0.0298,  0.0034]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.21.layer.0.SelfAttention.q.lora_B.weight': tensor([[-0.0021,  0.0006,  0.0097,  ..., -0.0011,  0.0026, -0.0374],\n",
+      "        [ 0.0022, -0.0127, -0.0116,  ...,  0.0035,  0.0007,  0.0061],\n",
+      "        [-0.0159,  0.0103,  0.0079,  ..., -0.0159, -0.0147, -0.0258],\n",
+      "        ...,\n",
+      "        [-0.0059,  0.0005, -0.0015,  ..., -0.0041, -0.0017, -0.0200],\n",
+      "        [-0.0226,  0.0153,  0.0165,  ..., -0.0202, -0.0170,  0.0057],\n",
+      "        [-0.0090,  0.0059,  0.0104,  ..., -0.0058, -0.0050, -0.0039]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.21.layer.0.SelfAttention.v.lora_A.weight': tensor([[-2.9747e-03,  1.7645e-05, -1.3959e-02,  ...,  1.6702e-02,\n",
+      "          2.5911e-02, -1.3416e-02],\n",
+      "        [ 1.5099e-02,  3.0057e-02, -2.4511e-02,  ..., -9.5109e-03,\n",
+      "         -1.4320e-02, -3.5647e-02],\n",
+      "        [ 8.3559e-03, -2.1729e-02,  6.0549e-03,  ...,  9.1296e-04,\n",
+      "          7.1130e-03,  6.5738e-03],\n",
+      "        ...,\n",
+      "        [ 2.6930e-02, -1.5346e-02,  8.9250e-03,  ...,  2.1033e-02,\n",
+      "         -1.1156e-02,  1.5283e-02],\n",
+      "        [ 3.1569e-02, -1.3637e-02, -4.1950e-02,  ..., -1.9484e-02,\n",
+      "          4.6394e-02, -1.6220e-02],\n",
+      "        [ 3.5849e-03,  3.0520e-02, -2.6064e-02,  ..., -2.2740e-02,\n",
+      "          4.4686e-02,  4.4240e-02]], device='cuda:0'), 'base_model.model.encoder.block.21.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0174, -0.0201, -0.0169,  ...,  0.0173,  0.0179,  0.0225],\n",
+      "        [ 0.0099, -0.0035, -0.0075,  ..., -0.0073, -0.0029, -0.0030],\n",
+      "        [ 0.0022, -0.0011, -0.0103,  ...,  0.0097,  0.0104,  0.0053],\n",
+      "        ...,\n",
+      "        [ 0.0152, -0.0175,  0.0037,  ...,  0.0109, -0.0065,  0.0123],\n",
+      "        [-0.0031, -0.0083,  0.0029,  ...,  0.0136, -0.0016, -0.0011],\n",
+      "        [-0.0010,  0.0170, -0.0216,  ..., -0.0187,  0.0204,  0.0015]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.22.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0314, -0.0045, -0.0207,  ..., -0.0543, -0.0126,  0.0270],\n",
+      "        [ 0.0462,  0.0351, -0.0017,  ...,  0.0194,  0.0221,  0.0235],\n",
+      "        [ 0.0180,  0.0013,  0.0002,  ..., -0.0235, -0.0701, -0.0042],\n",
+      "        ...,\n",
+      "        [-0.0203, -0.0024,  0.0114,  ..., -0.0052, -0.0093, -0.0337],\n",
+      "        [-0.0250,  0.0583,  0.0322,  ...,  0.0131,  0.0062, -0.0221],\n",
+      "        [-0.0365, -0.0355,  0.0022,  ..., -0.0091, -0.0091, -0.0029]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.22.layer.0.SelfAttention.q.lora_B.weight': tensor([[-0.0084,  0.0145, -0.0058,  ...,  0.0047,  0.0110, -0.0105],\n",
+      "        [ 0.0115,  0.0224,  0.0182,  ..., -0.0159, -0.0033, -0.0011],\n",
+      "        [-0.0078,  0.0098, -0.0079,  ...,  0.0062,  0.0060,  0.0024],\n",
+      "        ...,\n",
+      "        [ 0.0142, -0.0038,  0.0094,  ..., -0.0021, -0.0138,  0.0140],\n",
+      "        [ 0.0071, -0.0059, -0.0077,  ...,  0.0050, -0.0060,  0.0048],\n",
+      "        [ 0.0117, -0.0122, -0.0040,  ...,  0.0117, -0.0074,  0.0111]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.22.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0403,  0.0288,  0.0345,  ..., -0.0313,  0.0456,  0.0290],\n",
+      "        [ 0.0169, -0.0029, -0.0118,  ..., -0.0086,  0.0247, -0.0073],\n",
+      "        [-0.0182, -0.0154, -0.0023,  ...,  0.0213,  0.0123, -0.0090],\n",
+      "        ...,\n",
+      "        [ 0.0398, -0.0134, -0.0142,  ..., -0.0308,  0.0373,  0.0003],\n",
+      "        [-0.0289,  0.0193, -0.0009,  ...,  0.0243, -0.0158, -0.0300],\n",
+      "        [-0.0136, -0.0490, -0.0062,  ...,  0.0037, -0.0265, -0.0094]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.22.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 1.1087e-05, -2.0695e-02, -1.8504e-02,  ..., -2.3577e-02,\n",
+      "         -1.6880e-02,  2.1301e-02],\n",
+      "        [-2.5760e-02, -1.2344e-02, -8.9854e-03,  ...,  2.8194e-03,\n",
+      "         -1.9121e-02,  9.2568e-03],\n",
+      "        [-2.2456e-03, -1.0488e-02, -7.2482e-03,  ..., -1.6379e-02,\n",
+      "         -7.3226e-03,  1.2839e-02],\n",
+      "        ...,\n",
+      "        [-8.2132e-03, -1.3901e-02, -1.9611e-02,  ..., -1.4056e-02,\n",
+      "         -1.2747e-02,  1.2911e-02],\n",
+      "        [ 6.6918e-04, -9.2962e-03, -1.1142e-02,  ..., -8.8167e-03,\n",
+      "         -1.0819e-02,  1.0372e-02],\n",
+      "        [-6.8010e-03,  1.8507e-02,  2.1169e-02,  ...,  1.6902e-02,\n",
+      "          2.3338e-02, -1.5173e-02]], device='cuda:0'), 'base_model.model.encoder.block.23.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0287,  0.0413, -0.0387,  ...,  0.0297, -0.0248,  0.0058],\n",
+      "        [ 0.0030,  0.0372,  0.0231,  ..., -0.0069, -0.0422,  0.0422],\n",
+      "        [ 0.0030,  0.0138, -0.0341,  ...,  0.0274, -0.0051, -0.0047],\n",
+      "        ...,\n",
+      "        [-0.0029, -0.0164, -0.0410,  ...,  0.0089, -0.0333, -0.0027],\n",
+      "        [-0.0257, -0.0437,  0.0370,  ..., -0.0213,  0.0096, -0.0510],\n",
+      "        [-0.0032, -0.0417, -0.0039,  ..., -0.0076,  0.0019, -0.0326]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.23.layer.0.SelfAttention.q.lora_B.weight': tensor([[-0.0257, -0.0177, -0.0280,  ..., -0.0275,  0.0045,  0.0162],\n",
+      "        [ 0.0056,  0.0036,  0.0022,  ...,  0.0057, -0.0219, -0.0013],\n",
+      "        [ 0.0139,  0.0069,  0.0139,  ...,  0.0051,  0.0040, -0.0113],\n",
+      "        ...,\n",
+      "        [ 0.0008, -0.0016, -0.0095,  ...,  0.0090,  0.0071,  0.0025],\n",
+      "        [-0.0113, -0.0092, -0.0026,  ..., -0.0123, -0.0069,  0.0113],\n",
+      "        [ 0.0099,  0.0207,  0.0191,  ...,  0.0130, -0.0135, -0.0131]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.23.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0257, -0.0324,  0.0004,  ..., -0.0116,  0.0031,  0.0029],\n",
+      "        [-0.0007, -0.0106, -0.0271,  ..., -0.0073,  0.0254,  0.0318],\n",
+      "        [ 0.0251, -0.0242,  0.0158,  ...,  0.0124, -0.0365,  0.0444],\n",
+      "        ...,\n",
+      "        [ 0.0357,  0.0193,  0.0192,  ...,  0.0061, -0.0062, -0.0068],\n",
+      "        [-0.0283, -0.0168,  0.0212,  ...,  0.0222,  0.0310, -0.0504],\n",
+      "        [ 0.0458,  0.0061,  0.0074,  ..., -0.0495, -0.0247,  0.0176]],\n",
+      "       device='cuda:0'), 'base_model.model.encoder.block.23.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0123, -0.0024, -0.0170,  ..., -0.0145,  0.0109,  0.0031],\n",
+      "        [ 0.0129, -0.0127, -0.0100,  ..., -0.0079,  0.0134, -0.0099],\n",
+      "        [ 0.0170, -0.0094, -0.0102,  ..., -0.0114,  0.0088, -0.0183],\n",
+      "        ...,\n",
+      "        [ 0.0126, -0.0147, -0.0172,  ..., -0.0108,  0.0128, -0.0092],\n",
+      "        [ 0.0133, -0.0093, -0.0106,  ..., -0.0100,  0.0082, -0.0139],\n",
+      "        [ 0.0047, -0.0057, -0.0064,  ..., -0.0078,  0.0074,  0.0181]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.0.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0189, -0.0005,  0.0231,  ..., -0.0206, -0.0043, -0.0048],\n",
+      "        [-0.0072,  0.0131,  0.0362,  ...,  0.0150, -0.0252,  0.0271],\n",
+      "        [ 0.0177,  0.0022, -0.0043,  ...,  0.0060,  0.0091, -0.0020],\n",
+      "        ...,\n",
+      "        [-0.0144,  0.0091,  0.0296,  ..., -0.0143,  0.0180, -0.0202],\n",
+      "        [-0.0311,  0.0043,  0.0205,  ..., -0.0082, -0.0271,  0.0369],\n",
+      "        [ 0.0101,  0.0102, -0.0180,  ..., -0.0206,  0.0186,  0.0433]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.0.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0065, -0.0055,  0.0070,  ...,  0.0076, -0.0065, -0.0067],\n",
+      "        [-0.0073,  0.0057, -0.0076,  ..., -0.0072,  0.0075,  0.0079],\n",
+      "        [-0.0062,  0.0056, -0.0066,  ..., -0.0065,  0.0065,  0.0077],\n",
+      "        ...,\n",
+      "        [ 0.0048,  0.0183,  0.0174,  ..., -0.0067,  0.0044,  0.0170],\n",
+      "        [-0.0057,  0.0079,  0.0116,  ..., -0.0142,  0.0133,  0.0104],\n",
+      "        [-0.0040, -0.0191, -0.0187,  ...,  0.0082, -0.0056, -0.0184]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.0.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0238,  0.0030,  0.0195,  ..., -0.0004, -0.0226, -0.0122],\n",
+      "        [-0.0099, -0.0346, -0.0219,  ...,  0.0218,  0.0104, -0.0102],\n",
+      "        [-0.0086, -0.0257,  0.0187,  ...,  0.0245,  0.0112, -0.0370],\n",
+      "        ...,\n",
+      "        [ 0.0131, -0.0005, -0.0091,  ..., -0.0176, -0.0367, -0.0028],\n",
+      "        [ 0.0249,  0.0116,  0.0185,  ..., -0.0210,  0.0073,  0.0264],\n",
+      "        [-0.0072, -0.0039, -0.0250,  ...,  0.0153,  0.0267, -0.0015]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.0.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0131,  0.0152,  0.0148,  ...,  0.0143, -0.0135, -0.0138],\n",
+      "        [-0.0036,  0.0048,  0.0048,  ...,  0.0043, -0.0045, -0.0045],\n",
+      "        [-0.0127,  0.0133,  0.0084,  ...,  0.0093, -0.0077, -0.0085],\n",
+      "        ...,\n",
+      "        [-0.0257,  0.0265,  0.0275,  ...,  0.0272, -0.0283, -0.0266],\n",
+      "        [ 0.0154, -0.0107, -0.0156,  ..., -0.0120,  0.0129,  0.0101],\n",
+      "        [-0.0262,  0.0268,  0.0253,  ...,  0.0273, -0.0244, -0.0263]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.0.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0221, -0.0252,  0.0064,  ..., -0.0018, -0.0022, -0.0222],\n",
+      "        [-0.0064,  0.0094, -0.0091,  ...,  0.0109, -0.0069, -0.0229],\n",
+      "        [-0.0291,  0.0128,  0.0075,  ..., -0.0042, -0.0206,  0.0037],\n",
+      "        ...,\n",
+      "        [-0.0383, -0.0205,  0.0227,  ..., -0.0033, -0.0002,  0.0206],\n",
+      "        [ 0.0220,  0.0185, -0.0156,  ...,  0.0298,  0.0298, -0.0049],\n",
+      "        [ 0.0445, -0.0044, -0.0248,  ...,  0.0415, -0.0142,  0.0047]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.0.layer.1.EncDecAttention.q.lora_B.weight': tensor([[ 0.0167,  0.0058,  0.0135,  ...,  0.0126, -0.0094, -0.0168],\n",
+      "        [ 0.0015,  0.0061,  0.0088,  ...,  0.0045, -0.0071, -0.0033],\n",
+      "        [-0.0220,  0.0031, -0.0131,  ..., -0.0210,  0.0146,  0.0125],\n",
+      "        ...,\n",
+      "        [ 0.0155, -0.0182,  0.0175,  ...,  0.0183, -0.0085, -0.0135],\n",
+      "        [ 0.0049,  0.0088,  0.0103,  ...,  0.0133, -0.0007, -0.0076],\n",
+      "        [-0.0018, -0.0202, -0.0031,  ...,  0.0013,  0.0120,  0.0030]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.0.layer.1.EncDecAttention.v.lora_A.weight': tensor([[-0.0206, -0.0007,  0.0363,  ..., -0.0155,  0.0125,  0.0302],\n",
+      "        [ 0.0071, -0.0351, -0.0071,  ..., -0.0167,  0.0122, -0.0117],\n",
+      "        [ 0.0421, -0.0014, -0.0129,  ..., -0.0072, -0.0114, -0.0469],\n",
+      "        ...,\n",
+      "        [-0.0363, -0.0197,  0.0016,  ..., -0.0347, -0.0018,  0.0446],\n",
+      "        [ 0.0193,  0.0251, -0.0119,  ...,  0.0415,  0.0172, -0.0267],\n",
+      "        [ 0.0162,  0.0387, -0.0214,  ..., -0.0331,  0.0331, -0.0025]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.0.layer.1.EncDecAttention.v.lora_B.weight': tensor([[-0.0207,  0.0257,  0.0249,  ..., -0.0258,  0.0105, -0.0223],\n",
+      "        [-0.0098,  0.0059,  0.0131,  ..., -0.0124, -0.0172, -0.0154],\n",
+      "        [-0.0001, -0.0220,  0.0026,  ...,  0.0013, -0.0207,  0.0071],\n",
+      "        ...,\n",
+      "        [-0.0171,  0.0202,  0.0187,  ..., -0.0191, -0.0165, -0.0220],\n",
+      "        [ 0.0091, -0.0049, -0.0022,  ...,  0.0102,  0.0082,  0.0143],\n",
+      "        [ 0.0122, -0.0073, -0.0114,  ...,  0.0135, -0.0045,  0.0117]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.1.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0062,  0.0340, -0.0126,  ..., -0.0317, -0.0148,  0.0114],\n",
+      "        [ 0.0016,  0.0201, -0.0259,  ..., -0.0081, -0.0180,  0.0036],\n",
+      "        [-0.0196, -0.0027,  0.0157,  ..., -0.0040, -0.0164, -0.0303],\n",
+      "        ...,\n",
+      "        [-0.0307, -0.0401,  0.0096,  ..., -0.0108,  0.0102, -0.0025],\n",
+      "        [ 0.0146,  0.0296,  0.0074,  ...,  0.0368,  0.0319,  0.0163],\n",
+      "        [ 0.0212,  0.0060, -0.0247,  ...,  0.0213,  0.0096,  0.0179]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.1.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0083,  0.0084, -0.0099,  ...,  0.0071, -0.0083, -0.0079],\n",
+      "        [-0.0076, -0.0083,  0.0089,  ..., -0.0073,  0.0083,  0.0083],\n",
+      "        [-0.0078, -0.0082,  0.0091,  ..., -0.0071,  0.0082,  0.0078],\n",
+      "        ...,\n",
+      "        [ 0.0186, -0.0079, -0.0116,  ..., -0.0011,  0.0153,  0.0145],\n",
+      "        [ 0.0055, -0.0129,  0.0063,  ..., -0.0055, -0.0006,  0.0057],\n",
+      "        [-0.0068,  0.0094,  0.0049,  ..., -0.0214, -0.0105, -0.0143]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.1.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0069,  0.0235, -0.0279,  ...,  0.0227, -0.0086,  0.0170],\n",
+      "        [-0.0380, -0.0010,  0.0127,  ..., -0.0040, -0.0202, -0.0155],\n",
+      "        [ 0.0048, -0.0332,  0.0042,  ...,  0.0085, -0.0161, -0.0132],\n",
+      "        ...,\n",
+      "        [ 0.0162, -0.0120, -0.0343,  ...,  0.0348, -0.0002,  0.0094],\n",
+      "        [ 0.0057, -0.0085, -0.0151,  ..., -0.0225, -0.0477, -0.0247],\n",
+      "        [ 0.0267, -0.0135, -0.0142,  ...,  0.0340, -0.0055, -0.0230]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.1.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0185, -0.0192,  0.0192,  ...,  0.0191, -0.0146,  0.0189],\n",
+      "        [-0.0069, -0.0093,  0.0093,  ...,  0.0110, -0.0066,  0.0097],\n",
+      "        [-0.0158, -0.0179,  0.0180,  ...,  0.0171, -0.0151,  0.0170],\n",
+      "        ...,\n",
+      "        [ 0.0105,  0.0053, -0.0043,  ..., -0.0051,  0.0076, -0.0062],\n",
+      "        [-0.0061, -0.0023,  0.0024,  ...,  0.0026, -0.0050,  0.0019],\n",
+      "        [-0.0016, -0.0042,  0.0057,  ...,  0.0055,  0.0083,  0.0062]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.1.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0090,  0.0097,  0.0331,  ...,  0.0069, -0.0390, -0.0076],\n",
+      "        [-0.0125,  0.0190, -0.0041,  ...,  0.0240, -0.0379,  0.0208],\n",
+      "        [ 0.0228, -0.0280,  0.0044,  ...,  0.0122,  0.0094,  0.0115],\n",
+      "        ...,\n",
+      "        [-0.0252,  0.0296,  0.0183,  ..., -0.0186,  0.0372,  0.0183],\n",
+      "        [-0.0111, -0.0163,  0.0142,  ...,  0.0146, -0.0015, -0.0034],\n",
+      "        [ 0.0198,  0.0215, -0.0137,  ..., -0.0140,  0.0162,  0.0132]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.1.layer.1.EncDecAttention.q.lora_B.weight': tensor([[-0.0008,  0.0053,  0.0051,  ...,  0.0054, -0.0029,  0.0046],\n",
+      "        [ 0.0004,  0.0072,  0.0072,  ...,  0.0089,  0.0025,  0.0062],\n",
+      "        [ 0.0051, -0.0091, -0.0187,  ..., -0.0199,  0.0004, -0.0147],\n",
+      "        ...,\n",
+      "        [ 0.0021,  0.0219,  0.0079,  ...,  0.0080,  0.0092,  0.0068],\n",
+      "        [-0.0025, -0.0209, -0.0022,  ..., -0.0014, -0.0061, -0.0011],\n",
+      "        [ 0.0134,  0.0163, -0.0054,  ..., -0.0052,  0.0140, -0.0088]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.1.layer.1.EncDecAttention.v.lora_A.weight': tensor([[-0.0112,  0.0366, -0.0094,  ..., -0.0264,  0.0193,  0.0159],\n",
+      "        [-0.0172, -0.0171, -0.0273,  ...,  0.0062, -0.0069, -0.0036],\n",
+      "        [ 0.0194,  0.0024, -0.0140,  ..., -0.0079,  0.0263, -0.0028],\n",
+      "        ...,\n",
+      "        [ 0.0128, -0.0109, -0.0209,  ..., -0.0019, -0.0316, -0.0220],\n",
+      "        [ 0.0199,  0.0126,  0.0240,  ..., -0.0239,  0.0473,  0.0146],\n",
+      "        [ 0.0186,  0.0341, -0.0052,  ..., -0.0248,  0.0462,  0.0117]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.1.layer.1.EncDecAttention.v.lora_B.weight': tensor([[ 2.3198e-02, -2.1747e-02,  2.1701e-02,  ..., -2.1566e-02,\n",
+      "          2.4150e-02,  2.1774e-02],\n",
+      "        [-1.6284e-03,  7.0164e-04, -1.1504e-03,  ...,  4.7365e-03,\n",
+      "         -9.8117e-05, -6.4463e-03],\n",
+      "        [ 1.5113e-02, -1.2611e-02,  1.4666e-02,  ..., -9.9986e-03,\n",
+      "          1.1651e-02,  2.0736e-03],\n",
+      "        ...,\n",
+      "        [ 6.9344e-03, -3.2302e-03,  5.0286e-03,  ..., -7.6372e-03,\n",
+      "          3.3706e-04,  3.2491e-04],\n",
+      "        [ 1.6321e-02, -6.3958e-03,  5.0811e-03,  ..., -4.7124e-03,\n",
+      "          3.7874e-03,  3.6046e-03],\n",
+      "        [-9.5407e-04, -9.5503e-03,  1.0742e-02,  ..., -4.9108e-03,\n",
+      "          6.3167e-03,  5.7234e-03]], device='cuda:0'), 'base_model.model.decoder.block.2.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0043,  0.0264, -0.0115,  ...,  0.0311,  0.0349,  0.0248],\n",
+      "        [-0.0181, -0.0033,  0.0367,  ...,  0.0160, -0.0377, -0.0214],\n",
+      "        [ 0.0447, -0.0075, -0.0205,  ...,  0.0276,  0.0195, -0.0050],\n",
+      "        ...,\n",
+      "        [-0.0450, -0.0030,  0.0271,  ..., -0.0239,  0.0129,  0.0078],\n",
+      "        [-0.0393, -0.0002,  0.0071,  ..., -0.0083,  0.0064, -0.0282],\n",
+      "        [ 0.0125, -0.0222,  0.0274,  ..., -0.0385, -0.0269, -0.0489]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.2.layer.0.SelfAttention.q.lora_B.weight': tensor([[-3.1170e-03,  1.0720e-02,  2.1632e-03,  ..., -8.6142e-04,\n",
+      "         -1.2131e-03,  3.1180e-03],\n",
+      "        [ 6.1501e-03, -1.2021e-02,  3.3474e-07,  ..., -8.4009e-03,\n",
+      "         -3.4500e-03, -8.8669e-03],\n",
+      "        [-7.0843e-03,  5.7099e-04, -1.1848e-02,  ...,  2.5516e-03,\n",
+      "          5.1409e-03,  1.3616e-03],\n",
+      "        ...,\n",
+      "        [ 2.2997e-03,  5.7992e-03, -2.0466e-02,  ...,  1.0433e-02,\n",
+      "          1.6501e-02, -2.6774e-03],\n",
+      "        [-6.2623e-03,  4.4403e-04, -1.5141e-02,  ...,  6.5718e-03,\n",
+      "          1.8738e-02,  4.3951e-03],\n",
+      "        [-4.7032e-03,  5.6032e-03, -1.0028e-02,  ...,  9.1350e-03,\n",
+      "          8.2068e-03,  4.3547e-03]], device='cuda:0'), 'base_model.model.decoder.block.2.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0162, -0.0015,  0.0074,  ..., -0.0003,  0.0158, -0.0235],\n",
+      "        [-0.0018, -0.0206,  0.0025,  ...,  0.0132, -0.0139,  0.0327],\n",
+      "        [ 0.0019, -0.0290, -0.0073,  ..., -0.0219,  0.0109,  0.0003],\n",
+      "        ...,\n",
+      "        [ 0.0291, -0.0131, -0.0077,  ..., -0.0039, -0.0046, -0.0121],\n",
+      "        [ 0.0109, -0.0214,  0.0038,  ...,  0.0181, -0.0220, -0.0040],\n",
+      "        [-0.0069, -0.0031,  0.0367,  ..., -0.0086,  0.0165,  0.0042]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.2.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 2.4285e-02,  1.9491e-02,  2.0244e-02,  ...,  2.0564e-02,\n",
+      "         -2.1249e-02,  2.2550e-02],\n",
+      "        [ 1.6146e-02,  2.1073e-02,  2.0157e-02,  ...,  2.2195e-02,\n",
+      "         -1.9886e-02,  1.3432e-02],\n",
+      "        [-2.4051e-03, -3.1521e-03, -2.6053e-03,  ..., -8.9358e-03,\n",
+      "          2.8409e-03, -6.3286e-03],\n",
+      "        ...,\n",
+      "        [ 8.3711e-05, -5.0593e-03,  6.4664e-04,  ..., -3.6170e-03,\n",
+      "         -2.2688e-05, -3.6618e-03],\n",
+      "        [ 2.6448e-02,  2.5704e-02,  2.2131e-02,  ...,  2.7836e-02,\n",
+      "         -2.6099e-02,  1.0297e-02],\n",
+      "        [-2.4451e-02, -2.4363e-02, -2.2445e-02,  ..., -2.3449e-02,\n",
+      "          2.3975e-02, -2.5908e-02]], device='cuda:0'), 'base_model.model.decoder.block.2.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-3.2235e-02, -2.4453e-02, -1.8897e-03,  ..., -6.8799e-03,\n",
+      "         -4.4862e-02, -1.3606e-02],\n",
+      "        [ 2.0254e-02, -1.5936e-02,  4.4509e-03,  ...,  2.1429e-04,\n",
+      "         -1.0414e-02,  9.5917e-03],\n",
+      "        [-3.7509e-02, -7.7726e-03,  8.3584e-03,  ...,  8.5235e-05,\n",
+      "         -4.5792e-02,  1.7381e-02],\n",
+      "        ...,\n",
+      "        [ 1.3659e-02, -3.2639e-02,  1.3874e-02,  ..., -9.9682e-03,\n",
+      "          7.1122e-03, -1.3168e-02],\n",
+      "        [ 2.7758e-02,  1.7779e-02,  2.0833e-02,  ...,  9.7087e-03,\n",
+      "          4.2300e-02,  1.4540e-02],\n",
+      "        [-4.2461e-03,  3.1331e-03,  2.5784e-02,  ..., -3.8201e-02,\n",
+      "          6.7687e-03, -1.6261e-02]], device='cuda:0'), 'base_model.model.decoder.block.2.layer.1.EncDecAttention.q.lora_B.weight': tensor([[-0.0192, -0.0265, -0.0246,  ...,  0.0291,  0.0167, -0.0226],\n",
+      "        [ 0.0060,  0.0069,  0.0018,  ..., -0.0137, -0.0044,  0.0027],\n",
+      "        [ 0.0109,  0.0068,  0.0197,  ..., -0.0124, -0.0072,  0.0091],\n",
+      "        ...,\n",
+      "        [ 0.0187, -0.0130,  0.0149,  ..., -0.0143, -0.0120,  0.0074],\n",
+      "        [-0.0028,  0.0222,  0.0006,  ..., -0.0003,  0.0002,  0.0081],\n",
+      "        [-0.0265, -0.0019, -0.0296,  ...,  0.0252,  0.0172, -0.0095]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.2.layer.1.EncDecAttention.v.lora_A.weight': tensor([[-0.0110, -0.0249, -0.0031,  ..., -0.0298,  0.0260,  0.0516],\n",
+      "        [-0.0553,  0.0024,  0.0013,  ..., -0.0013,  0.0018,  0.0188],\n",
+      "        [-0.0096, -0.0039, -0.0053,  ..., -0.0048,  0.0163,  0.0174],\n",
+      "        ...,\n",
+      "        [-0.0267, -0.0287,  0.0545,  ...,  0.0190,  0.0132,  0.0087],\n",
+      "        [-0.0091,  0.0195,  0.0129,  ..., -0.0401,  0.0185,  0.0505],\n",
+      "        [ 0.0057, -0.0250,  0.0268,  ..., -0.0211,  0.0299,  0.0019]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.2.layer.1.EncDecAttention.v.lora_B.weight': tensor([[ 1.3512e-02,  1.0094e-02,  1.0749e-02,  ...,  1.3937e-02,\n",
+      "          1.4577e-02,  1.3163e-02],\n",
+      "        [-1.6808e-02, -1.2401e-02, -2.2549e-02,  ..., -1.5579e-02,\n",
+      "         -1.5503e-02, -1.6389e-02],\n",
+      "        [-6.7631e-03, -4.8791e-03, -9.9732e-03,  ..., -8.7470e-03,\n",
+      "         -7.3061e-03, -9.7309e-03],\n",
+      "        ...,\n",
+      "        [ 7.1581e-05, -2.2548e-03, -8.9037e-03,  ..., -6.1649e-03,\n",
+      "         -4.6464e-03, -6.3443e-03],\n",
+      "        [-2.1520e-02, -1.6354e-02, -2.1507e-02,  ..., -1.8473e-02,\n",
+      "         -2.1734e-02, -2.0071e-02],\n",
+      "        [-2.2858e-04, -1.2750e-03,  5.7743e-04,  ...,  1.1762e-03,\n",
+      "          2.3638e-03,  2.9686e-03]], device='cuda:0'), 'base_model.model.decoder.block.3.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0072, -0.0215,  0.0113,  ..., -0.0083, -0.0112, -0.0261],\n",
+      "        [ 0.0192,  0.0184,  0.0049,  ..., -0.0295,  0.0275, -0.0116],\n",
+      "        [ 0.0280,  0.0047, -0.0025,  ...,  0.0011, -0.0008,  0.0254],\n",
+      "        ...,\n",
+      "        [-0.0189,  0.0099, -0.0184,  ..., -0.0014,  0.0009,  0.0082],\n",
+      "        [-0.0026, -0.0022,  0.0208,  ..., -0.0045,  0.0279, -0.0247],\n",
+      "        [ 0.0204, -0.0134, -0.0076,  ...,  0.0101, -0.0260,  0.0311]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.3.layer.0.SelfAttention.q.lora_B.weight': tensor([[-0.0130,  0.0056,  0.0038,  ..., -0.0028, -0.0109,  0.0039],\n",
+      "        [ 0.0065, -0.0119,  0.0148,  ..., -0.0166,  0.0059,  0.0121],\n",
+      "        [-0.0192,  0.0150,  0.0005,  ...,  0.0006, -0.0156,  0.0008],\n",
+      "        ...,\n",
+      "        [-0.0164, -0.0180,  0.0115,  ..., -0.0031, -0.0211,  0.0231],\n",
+      "        [ 0.0160,  0.0132, -0.0068,  ...,  0.0005,  0.0191, -0.0242],\n",
+      "        [ 0.0048, -0.0069,  0.0201,  ..., -0.0211, -0.0003,  0.0144]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.3.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0154, -0.0190, -0.0139,  ..., -0.0068,  0.0060, -0.0097],\n",
+      "        [-0.0084,  0.0118, -0.0295,  ..., -0.0322, -0.0392, -0.0278],\n",
+      "        [ 0.0255, -0.0257,  0.0247,  ...,  0.0118, -0.0110,  0.0428],\n",
+      "        ...,\n",
+      "        [ 0.0037,  0.0066,  0.0163,  ...,  0.0142, -0.0087, -0.0071],\n",
+      "        [ 0.0321,  0.0076,  0.0286,  ..., -0.0035,  0.0288,  0.0387],\n",
+      "        [-0.0373, -0.0299, -0.0162,  ..., -0.0288, -0.0450, -0.0064]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.3.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0055,  0.0059, -0.0068,  ...,  0.0066, -0.0059,  0.0058],\n",
+      "        [ 0.0017,  0.0039, -0.0019,  ...,  0.0005,  0.0024, -0.0008],\n",
+      "        [-0.0046, -0.0104,  0.0097,  ..., -0.0098,  0.0082, -0.0097],\n",
+      "        ...,\n",
+      "        [-0.0010,  0.0006,  0.0019,  ...,  0.0017, -0.0012, -0.0005],\n",
+      "        [ 0.0027,  0.0070, -0.0015,  ...,  0.0019, -0.0003,  0.0019],\n",
+      "        [-0.0111, -0.0121,  0.0130,  ..., -0.0129,  0.0124, -0.0135]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.3.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0305,  0.0111,  0.0111,  ...,  0.0119, -0.0292, -0.0201],\n",
+      "        [ 0.0215,  0.0097, -0.0449,  ..., -0.0125,  0.0371,  0.0259],\n",
+      "        [-0.0006, -0.0228,  0.0173,  ..., -0.0298, -0.0348, -0.0166],\n",
+      "        ...,\n",
+      "        [-0.0496, -0.0060,  0.0336,  ..., -0.0376, -0.0366,  0.0158],\n",
+      "        [ 0.0150, -0.0011, -0.0006,  ...,  0.0357,  0.0258,  0.0381],\n",
+      "        [-0.0085, -0.0071,  0.0198,  ...,  0.0362, -0.0071,  0.0058]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.3.layer.1.EncDecAttention.q.lora_B.weight': tensor([[-0.0022, -0.0042,  0.0066,  ...,  0.0047, -0.0036, -0.0020],\n",
+      "        [ 0.0189, -0.0371,  0.0382,  ...,  0.0305, -0.0319, -0.0206],\n",
+      "        [-0.0052, -0.0098,  0.0078,  ...,  0.0078, -0.0058, -0.0184],\n",
+      "        ...,\n",
+      "        [ 0.0121,  0.0118, -0.0107,  ..., -0.0072,  0.0152, -0.0095],\n",
+      "        [-0.0156,  0.0325, -0.0318,  ..., -0.0103,  0.0305,  0.0194],\n",
+      "        [ 0.0130,  0.0154, -0.0188,  ..., -0.0170,  0.0142,  0.0074]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.3.layer.1.EncDecAttention.v.lora_A.weight': tensor([[ 0.0188, -0.0262,  0.0216,  ...,  0.0069, -0.0215, -0.0048],\n",
+      "        [ 0.0322, -0.0194,  0.0094,  ...,  0.0157, -0.0021,  0.0539],\n",
+      "        [-0.0159,  0.0084,  0.0091,  ..., -0.0113,  0.0037,  0.0314],\n",
+      "        ...,\n",
+      "        [ 0.0004, -0.0198,  0.0015,  ..., -0.0551,  0.0388,  0.0306],\n",
+      "        [ 0.0350, -0.0306, -0.0190,  ...,  0.0378, -0.0301, -0.0262],\n",
+      "        [ 0.0203,  0.0142, -0.0191,  ..., -0.0139,  0.0464,  0.0181]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.3.layer.1.EncDecAttention.v.lora_B.weight': tensor([[-0.0019, -0.0041,  0.0024,  ..., -0.0003, -0.0006,  0.0029],\n",
+      "        [-0.0188,  0.0186,  0.0175,  ...,  0.0169, -0.0168,  0.0166],\n",
+      "        [ 0.0126, -0.0167, -0.0079,  ..., -0.0090,  0.0082, -0.0114],\n",
+      "        ...,\n",
+      "        [-0.0065,  0.0139,  0.0076,  ...,  0.0082, -0.0083,  0.0121],\n",
+      "        [-0.0111,  0.0149,  0.0127,  ...,  0.0131, -0.0134,  0.0118],\n",
+      "        [ 0.0017,  0.0036, -0.0023,  ..., -0.0028,  0.0031, -0.0028]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.4.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0342, -0.0317, -0.0014,  ..., -0.0271,  0.0120,  0.0191],\n",
+      "        [-0.0287, -0.0342,  0.0428,  ...,  0.0099, -0.0305,  0.0277],\n",
+      "        [-0.0308, -0.0273, -0.0005,  ..., -0.0049, -0.0236, -0.0056],\n",
+      "        ...,\n",
+      "        [-0.0138, -0.0167, -0.0405,  ...,  0.0315,  0.0343,  0.0241],\n",
+      "        [-0.0294, -0.0314,  0.0403,  ...,  0.0076, -0.0130, -0.0139],\n",
+      "        [-0.0245, -0.0102,  0.0352,  ..., -0.0369, -0.0388,  0.0226]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.4.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0039,  0.0033,  0.0120,  ...,  0.0007, -0.0020,  0.0129],\n",
+      "        [ 0.0154,  0.0150,  0.0096,  ..., -0.0125,  0.0154,  0.0157],\n",
+      "        [ 0.0114,  0.0126,  0.0129,  ..., -0.0122,  0.0115,  0.0021],\n",
+      "        ...,\n",
+      "        [-0.0076, -0.0078, -0.0041,  ..., -0.0005, -0.0058, -0.0095],\n",
+      "        [ 0.0091,  0.0096,  0.0089,  ..., -0.0087,  0.0105,  0.0073],\n",
+      "        [-0.0055, -0.0083, -0.0074,  ...,  0.0106, -0.0107, -0.0063]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.4.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0116,  0.0151, -0.0338,  ...,  0.0024,  0.0069,  0.0019],\n",
+      "        [-0.0075,  0.0279, -0.0158,  ..., -0.0128,  0.0325,  0.0015],\n",
+      "        [ 0.0325, -0.0339,  0.0060,  ..., -0.0077,  0.0005,  0.0309],\n",
+      "        ...,\n",
+      "        [ 0.0232,  0.0226, -0.0187,  ...,  0.0152, -0.0232,  0.0162],\n",
+      "        [ 0.0283,  0.0220,  0.0349,  ...,  0.0378,  0.0139,  0.0358],\n",
+      "        [ 0.0147,  0.0302, -0.0180,  ...,  0.0072, -0.0060, -0.0410]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.4.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0115,  0.0098,  0.0073,  ...,  0.0100,  0.0085, -0.0100],\n",
+      "        [ 0.0131, -0.0123, -0.0133,  ..., -0.0141, -0.0025,  0.0153],\n",
+      "        [-0.0161,  0.0144,  0.0136,  ...,  0.0137,  0.0108, -0.0134],\n",
+      "        ...,\n",
+      "        [-0.0047,  0.0031,  0.0027,  ...,  0.0071,  0.0149, -0.0057],\n",
+      "        [ 0.0156, -0.0158, -0.0076,  ..., -0.0147, -0.0157,  0.0160],\n",
+      "        [-0.0158,  0.0169,  0.0219,  ...,  0.0165,  0.0194, -0.0162]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.4.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0263,  0.0215,  0.0172,  ..., -0.0139, -0.0270, -0.0057],\n",
+      "        [ 0.0239, -0.0136,  0.0008,  ...,  0.0056,  0.0047, -0.0103],\n",
+      "        [ 0.0140,  0.0015, -0.0131,  ...,  0.0060,  0.0066,  0.0215],\n",
+      "        ...,\n",
+      "        [ 0.0302, -0.0049, -0.0103,  ...,  0.0099, -0.0047, -0.0163],\n",
+      "        [ 0.0014,  0.0275,  0.0259,  ...,  0.0201, -0.0150,  0.0098],\n",
+      "        [ 0.0081,  0.0232, -0.0162,  ..., -0.0027, -0.0194,  0.0108]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.4.layer.1.EncDecAttention.q.lora_B.weight': tensor([[ 0.0242,  0.0186, -0.0200,  ...,  0.0223, -0.0155,  0.0179],\n",
+      "        [-0.0203, -0.0142,  0.0140,  ..., -0.0150,  0.0092, -0.0160],\n",
+      "        [-0.0007,  0.0006, -0.0005,  ...,  0.0006, -0.0002, -0.0026],\n",
+      "        ...,\n",
+      "        [ 0.0092,  0.0041, -0.0047,  ...,  0.0082, -0.0055,  0.0086],\n",
+      "        [ 0.0151,  0.0137, -0.0022,  ...,  0.0107, -0.0017,  0.0091],\n",
+      "        [-0.0118, -0.0104,  0.0048,  ..., -0.0077, -0.0013, -0.0059]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.4.layer.1.EncDecAttention.v.lora_A.weight': tensor([[-0.0023, -0.0160,  0.0328,  ...,  0.0381,  0.0126, -0.0562],\n",
+      "        [-0.0314,  0.0148,  0.0142,  ...,  0.0375, -0.0140, -0.0295],\n",
+      "        [-0.0111,  0.0272,  0.0154,  ...,  0.0156,  0.0067, -0.0610],\n",
+      "        ...,\n",
+      "        [-0.0016, -0.0208,  0.0008,  ...,  0.0163,  0.0187,  0.0313],\n",
+      "        [ 0.0035, -0.0200,  0.0226,  ..., -0.0011, -0.0050,  0.0417],\n",
+      "        [ 0.0164,  0.0416, -0.0165,  ..., -0.0197,  0.0283, -0.0016]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.4.layer.1.EncDecAttention.v.lora_B.weight': tensor([[ 0.0080,  0.0172,  0.0179,  ...,  0.0031, -0.0199, -0.0094],\n",
+      "        [ 0.0060,  0.0121,  0.0121,  ..., -0.0026, -0.0113, -0.0111],\n",
+      "        [-0.0091, -0.0273, -0.0263,  ...,  0.0116,  0.0272,  0.0177],\n",
+      "        ...,\n",
+      "        [-0.0019, -0.0016,  0.0025,  ...,  0.0041,  0.0043,  0.0018],\n",
+      "        [-0.0042, -0.0017, -0.0011,  ..., -0.0032, -0.0010,  0.0023],\n",
+      "        [ 0.0004,  0.0092,  0.0103,  ..., -0.0058, -0.0173, -0.0072]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.5.layer.0.SelfAttention.q.lora_A.weight': tensor([[-2.6730e-02,  2.0864e-02,  2.3388e-02,  ..., -1.5293e-02,\n",
+      "         -1.6101e-03, -4.0400e-02],\n",
+      "        [-3.2931e-02, -1.0590e-03,  1.2445e-02,  ..., -5.5235e-03,\n",
+      "          3.3820e-02, -2.0991e-02],\n",
+      "        [ 1.2719e-02,  2.5506e-03, -3.6293e-02,  ...,  1.6096e-02,\n",
+      "          1.3012e-02,  1.6014e-02],\n",
+      "        ...,\n",
+      "        [-1.3189e-02, -2.4657e-02, -8.4241e-03,  ..., -5.5609e-03,\n",
+      "         -4.4067e-03, -3.3457e-02],\n",
+      "        [ 2.7041e-02,  2.7457e-02, -2.4774e-03,  ...,  2.2314e-03,\n",
+      "          1.1362e-02, -9.3929e-03],\n",
+      "        [ 1.5837e-02, -9.2169e-03, -7.3251e-04,  ..., -2.6933e-05,\n",
+      "          1.2411e-02, -4.1066e-03]], device='cuda:0'), 'base_model.model.decoder.block.5.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0014,  0.0082,  0.0049,  ..., -0.0001, -0.0062,  0.0002],\n",
+      "        [-0.0049, -0.0086, -0.0027,  ..., -0.0024,  0.0020,  0.0039],\n",
+      "        [-0.0048, -0.0091, -0.0022,  ..., -0.0030,  0.0021,  0.0041],\n",
+      "        ...,\n",
+      "        [-0.0004,  0.0162,  0.0006,  ...,  0.0005, -0.0079,  0.0013],\n",
+      "        [-0.0036, -0.0084,  0.0063,  ..., -0.0062,  0.0043,  0.0085],\n",
+      "        [-0.0044,  0.0027,  0.0058,  ..., -0.0065, -0.0014,  0.0068]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.5.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0186, -0.0075,  0.0282,  ..., -0.0128, -0.0348, -0.0162],\n",
+      "        [ 0.0120,  0.0150, -0.0364,  ..., -0.0095,  0.0025, -0.0103],\n",
+      "        [ 0.0250,  0.0160, -0.0067,  ..., -0.0135,  0.0271,  0.0212],\n",
+      "        ...,\n",
+      "        [ 0.0176, -0.0096,  0.0075,  ...,  0.0166, -0.0108, -0.0156],\n",
+      "        [ 0.0347,  0.0249, -0.0111,  ..., -0.0084,  0.0292,  0.0088],\n",
+      "        [-0.0105,  0.0333,  0.0109,  ..., -0.0365,  0.0088, -0.0124]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.5.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0106,  0.0103,  0.0088,  ...,  0.0104,  0.0139, -0.0114],\n",
+      "        [ 0.0053, -0.0025, -0.0100,  ..., -0.0028, -0.0017,  0.0061],\n",
+      "        [ 0.0084, -0.0051, -0.0125,  ..., -0.0073, -0.0063,  0.0097],\n",
+      "        ...,\n",
+      "        [ 0.0039,  0.0021, -0.0035,  ..., -0.0035, -0.0005,  0.0048],\n",
+      "        [-0.0107,  0.0107,  0.0084,  ...,  0.0114,  0.0113, -0.0080],\n",
+      "        [ 0.0057, -0.0105, -0.0100,  ..., -0.0036, -0.0060,  0.0014]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.5.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0253,  0.0348,  0.0144,  ...,  0.0535, -0.0086, -0.0021],\n",
+      "        [ 0.0331, -0.0260,  0.0130,  ..., -0.0028, -0.0292,  0.0152],\n",
+      "        [ 0.0248, -0.0087, -0.0264,  ..., -0.0248, -0.0025,  0.0110],\n",
+      "        ...,\n",
+      "        [ 0.0223, -0.0317,  0.0493,  ..., -0.0156, -0.0037, -0.0297],\n",
+      "        [ 0.0025,  0.0178, -0.0461,  ..., -0.0236,  0.0002,  0.0295],\n",
+      "        [ 0.0259, -0.0477,  0.0526,  ..., -0.0363, -0.0160, -0.0426]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.5.layer.1.EncDecAttention.q.lora_B.weight': tensor([[ 0.0095, -0.0096, -0.0092,  ..., -0.0071,  0.0069, -0.0087],\n",
+      "        [ 0.0015, -0.0024, -0.0026,  ..., -0.0026,  0.0081, -0.0055],\n",
+      "        [-0.0016,  0.0003,  0.0006,  ...,  0.0035, -0.0067, -0.0027],\n",
+      "        ...,\n",
+      "        [ 0.0003,  0.0031, -0.0029,  ...,  0.0003,  0.0018, -0.0005],\n",
+      "        [-0.0036,  0.0015,  0.0029,  ...,  0.0102, -0.0047,  0.0026],\n",
+      "        [-0.0088,  0.0082,  0.0067,  ..., -0.0148,  0.0157, -0.0170]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.5.layer.1.EncDecAttention.v.lora_A.weight': tensor([[-0.0108, -0.0207, -0.0201,  ..., -0.0062, -0.0595, -0.0536],\n",
+      "        [ 0.0466,  0.0231,  0.0192,  ...,  0.0111,  0.0117,  0.0167],\n",
+      "        [ 0.0248, -0.0007, -0.0118,  ..., -0.0208,  0.0139,  0.0163],\n",
+      "        ...,\n",
+      "        [ 0.0050, -0.0228,  0.0266,  ...,  0.0138,  0.0375,  0.0285],\n",
+      "        [ 0.0072,  0.0144,  0.0285,  ...,  0.0107,  0.0474,  0.0207],\n",
+      "        [ 0.0098,  0.0003, -0.0225,  ..., -0.0130,  0.0355,  0.0492]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.5.layer.1.EncDecAttention.v.lora_B.weight': tensor([[-2.1625e-03,  1.6814e-04,  8.1146e-05,  ...,  1.4400e-03,\n",
+      "         -5.4853e-03, -1.1964e-03],\n",
+      "        [ 6.4509e-03, -6.9242e-03, -1.6537e-02,  ..., -6.6245e-03,\n",
+      "         -2.1439e-02, -1.5015e-02],\n",
+      "        [-8.9251e-04,  7.1215e-04, -2.7612e-04,  ...,  2.6042e-04,\n",
+      "          1.3185e-02,  9.8170e-03],\n",
+      "        ...,\n",
+      "        [-3.6989e-03,  1.8064e-03,  8.6825e-03,  ...,  2.5273e-03,\n",
+      "          4.6486e-03,  5.4987e-03],\n",
+      "        [-1.1506e-02,  8.4292e-03,  5.8922e-03,  ...,  1.1264e-02,\n",
+      "          4.8699e-03,  7.4301e-03],\n",
+      "        [ 1.0491e-02, -9.8668e-03, -1.0596e-02,  ..., -9.7787e-03,\n",
+      "         -5.6693e-03, -9.9868e-03]], device='cuda:0'), 'base_model.model.decoder.block.6.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0067, -0.0038,  0.0387,  ...,  0.0008, -0.0110,  0.0104],\n",
+      "        [ 0.0157, -0.0423,  0.0041,  ..., -0.0276,  0.0129,  0.0129],\n",
+      "        [-0.0351,  0.0106, -0.0120,  ...,  0.0130, -0.0018, -0.0170],\n",
+      "        ...,\n",
+      "        [ 0.0046,  0.0349,  0.0031,  ...,  0.0075,  0.0278,  0.0109],\n",
+      "        [ 0.0355,  0.0083,  0.0434,  ..., -0.0034, -0.0004, -0.0221],\n",
+      "        [-0.0067,  0.0074,  0.0048,  ...,  0.0083, -0.0431, -0.0202]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.6.layer.0.SelfAttention.q.lora_B.weight': tensor([[-1.4334e-02, -2.0122e-02,  2.5548e-02,  ...,  1.8705e-02,\n",
+      "          4.1224e-04, -9.1952e-03],\n",
+      "        [-3.4477e-03, -6.8384e-04, -4.9736e-05,  ...,  2.7096e-03,\n",
+      "          9.1590e-03,  1.4286e-02],\n",
+      "        [-1.1760e-02, -1.6113e-02,  1.9543e-02,  ...,  1.4789e-02,\n",
+      "          3.0643e-03, -6.1182e-03],\n",
+      "        ...,\n",
+      "        [ 1.0866e-02,  1.0069e-02, -8.9908e-03,  ..., -9.9777e-03,\n",
+      "          1.0297e-02,  7.7944e-03],\n",
+      "        [-1.1795e-02, -1.1093e-02,  1.2621e-02,  ...,  1.2965e-02,\n",
+      "         -9.7311e-03, -1.1893e-02],\n",
+      "        [ 1.2635e-02,  1.1360e-02, -1.1950e-02,  ..., -1.2414e-02,\n",
+      "          1.0531e-02,  1.4034e-02]], device='cuda:0'), 'base_model.model.decoder.block.6.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0033, -0.0322,  0.0411,  ..., -0.0003, -0.0140,  0.0132],\n",
+      "        [ 0.0124, -0.0387,  0.0131,  ..., -0.0345,  0.0151,  0.0003],\n",
+      "        [-0.0095,  0.0191, -0.0311,  ..., -0.0108, -0.0286, -0.0093],\n",
+      "        ...,\n",
+      "        [-0.0186,  0.0239, -0.0576,  ..., -0.0223, -0.0066, -0.0108],\n",
+      "        [ 0.0070, -0.0496, -0.0148,  ..., -0.0102, -0.0159,  0.0199],\n",
+      "        [ 0.0071,  0.0019, -0.0213,  ...,  0.0124, -0.0183,  0.0216]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.6.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 2.4060e-02,  2.1683e-02, -2.5614e-02,  ..., -2.4639e-02,\n",
+      "          2.5046e-02,  2.4352e-02],\n",
+      "        [ 1.3227e-02,  9.1044e-03, -1.3184e-02,  ..., -1.1382e-02,\n",
+      "          1.2999e-02,  1.3374e-02],\n",
+      "        [-1.4519e-02, -1.2697e-02,  1.2229e-02,  ...,  1.1532e-02,\n",
+      "         -1.4862e-02, -1.5476e-02],\n",
+      "        ...,\n",
+      "        [-1.6640e-02, -1.5588e-02,  1.2420e-02,  ...,  2.0494e-02,\n",
+      "         -1.3000e-02, -1.1540e-02],\n",
+      "        [ 1.6104e-02,  1.6665e-02, -8.1493e-03,  ..., -2.1218e-02,\n",
+      "          8.9385e-03,  8.0249e-03],\n",
+      "        [ 6.4023e-05,  1.2668e-03, -8.7618e-03,  ...,  3.2057e-05,\n",
+      "          8.1570e-03,  9.9918e-03]], device='cuda:0'), 'base_model.model.decoder.block.6.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0133,  0.0064,  0.0368,  ...,  0.0051, -0.0354, -0.0184],\n",
+      "        [-0.0097,  0.0049, -0.0109,  ..., -0.0311, -0.0086, -0.0254],\n",
+      "        [-0.0274,  0.0319, -0.0215,  ...,  0.0142,  0.0177,  0.0154],\n",
+      "        ...,\n",
+      "        [-0.0239, -0.0018,  0.0185,  ..., -0.0107, -0.0099,  0.0124],\n",
+      "        [-0.0030, -0.0054, -0.0009,  ..., -0.0291,  0.0086,  0.0105],\n",
+      "        [-0.0358, -0.0142, -0.0217,  ..., -0.0015,  0.0169,  0.0144]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.6.layer.1.EncDecAttention.q.lora_B.weight': tensor([[ 0.0037,  0.0140,  0.0248,  ...,  0.0287, -0.0060,  0.0007],\n",
+      "        [-0.0074,  0.0020, -0.0056,  ...,  0.0008,  0.0032,  0.0039],\n",
+      "        [-0.0019, -0.0019,  0.0006,  ...,  0.0045, -0.0034,  0.0044],\n",
+      "        ...,\n",
+      "        [ 0.0046, -0.0029,  0.0052,  ...,  0.0034,  0.0060, -0.0008],\n",
+      "        [ 0.0097,  0.0037,  0.0133,  ...,  0.0158, -0.0098, -0.0095],\n",
+      "        [ 0.0096, -0.0046,  0.0003,  ..., -0.0082,  0.0195, -0.0122]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.6.layer.1.EncDecAttention.v.lora_A.weight': tensor([[ 0.0266, -0.0150, -0.0428,  ...,  0.0195,  0.0117,  0.0346],\n",
+      "        [ 0.0236, -0.0423,  0.0263,  ..., -0.0011, -0.0402, -0.0415],\n",
+      "        [-0.0090,  0.0466, -0.0165,  ..., -0.0530, -0.0045,  0.0174],\n",
+      "        ...,\n",
+      "        [-0.0144, -0.0242,  0.0267,  ...,  0.0429, -0.0238, -0.0078],\n",
+      "        [ 0.0131, -0.0153,  0.0209,  ...,  0.0367, -0.0131, -0.0385],\n",
+      "        [ 0.0002,  0.0393, -0.0139,  ..., -0.0045,  0.0177,  0.0329]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.6.layer.1.EncDecAttention.v.lora_B.weight': tensor([[-0.0009, -0.0140,  0.0066,  ..., -0.0093,  0.0028,  0.0081],\n",
+      "        [-0.0098,  0.0039, -0.0037,  ...,  0.0026,  0.0042, -0.0022],\n",
+      "        [ 0.0084,  0.0001,  0.0004,  ...,  0.0047, -0.0066,  0.0003],\n",
+      "        ...,\n",
+      "        [ 0.0044,  0.0213, -0.0220,  ...,  0.0198,  0.0182, -0.0189],\n",
+      "        [-0.0040,  0.0051, -0.0058,  ...,  0.0041,  0.0030, -0.0046],\n",
+      "        [-0.0051,  0.0010,  0.0045,  ..., -0.0046, -0.0037,  0.0040]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.7.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0331, -0.0111, -0.0323,  ...,  0.0251,  0.0505, -0.0046],\n",
+      "        [-0.0215,  0.0399,  0.0451,  ..., -0.0216, -0.0509, -0.0254],\n",
+      "        [-0.0085,  0.0070,  0.0028,  ..., -0.0233,  0.0098, -0.0006],\n",
+      "        ...,\n",
+      "        [-0.0136, -0.0158,  0.0197,  ..., -0.0471,  0.0237, -0.0038],\n",
+      "        [ 0.0197, -0.0128, -0.0333,  ...,  0.0287,  0.0450, -0.0127],\n",
+      "        [ 0.0028,  0.0185, -0.0092,  ..., -0.0441, -0.0428,  0.0232]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.7.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0008,  0.0002, -0.0002,  ..., -0.0006,  0.0024,  0.0010],\n",
+      "        [ 0.0004,  0.0005, -0.0004,  ..., -0.0005,  0.0018,  0.0012],\n",
+      "        [ 0.0012, -0.0009,  0.0010,  ..., -0.0017,  0.0024, -0.0002],\n",
+      "        ...,\n",
+      "        [ 0.0071, -0.0029, -0.0035,  ..., -0.0041,  0.0080, -0.0040],\n",
+      "        [ 0.0006, -0.0016,  0.0108,  ...,  0.0008, -0.0043, -0.0014],\n",
+      "        [ 0.0016,  0.0023,  0.0044,  ...,  0.0002,  0.0046,  0.0030]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.7.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0091, -0.0108, -0.0018,  ...,  0.0179,  0.0270,  0.0174],\n",
+      "        [-0.0182,  0.0060, -0.0365,  ...,  0.0066, -0.0001, -0.0285],\n",
+      "        [-0.0408,  0.0361, -0.0176,  ...,  0.0240,  0.0034,  0.0067],\n",
+      "        ...,\n",
+      "        [ 0.0293,  0.0057,  0.0186,  ...,  0.0305, -0.0190,  0.0132],\n",
+      "        [-0.0201, -0.0120,  0.0187,  ..., -0.0276, -0.0281, -0.0402],\n",
+      "        [ 0.0082, -0.0213,  0.0019,  ...,  0.0298, -0.0016, -0.0259]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.7.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0195, -0.0055, -0.0053,  ...,  0.0047, -0.0101,  0.0149],\n",
+      "        [ 0.0014, -0.0081, -0.0072,  ...,  0.0070, -0.0058, -0.0003],\n",
+      "        [ 0.0088,  0.0007,  0.0048,  ..., -0.0033,  0.0032,  0.0156],\n",
+      "        ...,\n",
+      "        [ 0.0045, -0.0026, -0.0015,  ...,  0.0017, -0.0012,  0.0068],\n",
+      "        [ 0.0081, -0.0016, -0.0009,  ...,  0.0013, -0.0054,  0.0117],\n",
+      "        [ 0.0039, -0.0103, -0.0033,  ...,  0.0092, -0.0028,  0.0041]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.7.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0203,  0.0043,  0.0081,  ...,  0.0358,  0.0223,  0.0172],\n",
+      "        [-0.0086, -0.0177,  0.0020,  ..., -0.0164, -0.0001,  0.0288],\n",
+      "        [-0.0135,  0.0085,  0.0117,  ...,  0.0381,  0.0329,  0.0229],\n",
+      "        ...,\n",
+      "        [-0.0364, -0.0218, -0.0393,  ...,  0.0238,  0.0459,  0.0361],\n",
+      "        [ 0.0031, -0.0134, -0.0096,  ...,  0.0044, -0.0182, -0.0195],\n",
+      "        [ 0.0105,  0.0036, -0.0233,  ...,  0.0360,  0.0076,  0.0006]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.7.layer.1.EncDecAttention.q.lora_B.weight': tensor([[-0.0246, -0.0249,  0.0212,  ..., -0.0247,  0.0232, -0.0254],\n",
+      "        [-0.0233, -0.0259,  0.0089,  ..., -0.0237,  0.0226, -0.0243],\n",
+      "        [-0.0282, -0.0282,  0.0009,  ..., -0.0280,  0.0279, -0.0301],\n",
+      "        ...,\n",
+      "        [-0.0124, -0.0138,  0.0094,  ..., -0.0132,  0.0121, -0.0126],\n",
+      "        [ 0.0109,  0.0019, -0.0137,  ...,  0.0081, -0.0119,  0.0121],\n",
+      "        [ 0.0128,  0.0076,  0.0082,  ...,  0.0131, -0.0137,  0.0123]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.7.layer.1.EncDecAttention.v.lora_A.weight': tensor([[-0.0427,  0.0195,  0.0449,  ...,  0.0135, -0.0402, -0.0507],\n",
+      "        [ 0.0125,  0.0196, -0.0151,  ..., -0.0420,  0.0035,  0.0245],\n",
+      "        [ 0.0263, -0.0365,  0.0320,  ...,  0.0032, -0.0379,  0.0087],\n",
+      "        ...,\n",
+      "        [ 0.0331,  0.0281, -0.0341,  ...,  0.0250,  0.0080,  0.0225],\n",
+      "        [-0.0033, -0.0399, -0.0354,  ...,  0.0027,  0.0030,  0.0226],\n",
+      "        [-0.0005,  0.0128,  0.0210,  ..., -0.0145,  0.0167, -0.0309]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.7.layer.1.EncDecAttention.v.lora_B.weight': tensor([[ 8.7193e-03, -9.4752e-03,  4.8379e-03,  ..., -6.9801e-03,\n",
+      "          4.6064e-03,  1.4173e-02],\n",
+      "        [ 1.0323e-02, -7.7941e-03,  1.0070e-02,  ..., -2.5394e-03,\n",
+      "          7.4860e-03, -1.9913e-03],\n",
+      "        [ 2.3713e-03, -3.8002e-03,  6.1123e-03,  ..., -6.1484e-03,\n",
+      "          8.3422e-03, -1.3110e-02],\n",
+      "        ...,\n",
+      "        [ 2.4322e-03,  1.1196e-02, -8.1483e-03,  ..., -9.3193e-03,\n",
+      "         -4.1116e-03,  6.0736e-03],\n",
+      "        [-4.6512e-03,  2.2509e-05,  2.8143e-03,  ...,  7.7112e-03,\n",
+      "          1.1012e-02, -1.5901e-02],\n",
+      "        [-2.5353e-02,  1.9855e-02, -1.2999e-02,  ...,  2.1589e-02,\n",
+      "          3.5988e-03, -1.5793e-02]], device='cuda:0'), 'base_model.model.decoder.block.8.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0055,  0.0482, -0.0049,  ...,  0.0258,  0.0268,  0.0282],\n",
+      "        [ 0.0004,  0.0061,  0.0378,  ..., -0.0061,  0.0131, -0.0106],\n",
+      "        [-0.0127, -0.0368,  0.0114,  ..., -0.0245, -0.0108,  0.0082],\n",
+      "        ...,\n",
+      "        [ 0.0270, -0.0469,  0.0392,  ..., -0.0057,  0.0246, -0.0449],\n",
+      "        [ 0.0149, -0.0116,  0.0343,  ...,  0.0086, -0.0134, -0.0224],\n",
+      "        [-0.0031, -0.0186,  0.0189,  ...,  0.0102,  0.0251,  0.0056]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.8.layer.0.SelfAttention.q.lora_B.weight': tensor([[-8.1862e-03,  1.1881e-02,  9.7900e-03,  ...,  1.0970e-02,\n",
+      "         -4.8307e-03,  1.3177e-02],\n",
+      "        [ 1.3475e-02, -1.3164e-02, -1.6153e-02,  ..., -3.9508e-03,\n",
+      "         -5.1986e-05, -6.6920e-03],\n",
+      "        [ 1.3827e-03,  3.2577e-03, -4.9111e-03,  ...,  5.6041e-03,\n",
+      "          3.4114e-03, -1.3535e-03],\n",
+      "        ...,\n",
+      "        [-2.6834e-02,  2.3956e-02,  2.7925e-02,  ...,  2.6604e-02,\n",
+      "          2.5881e-02, -4.2479e-03],\n",
+      "        [ 9.8006e-03,  5.8498e-03, -1.5935e-02,  ..., -1.2355e-02,\n",
+      "         -2.2700e-03,  1.5965e-02],\n",
+      "        [ 1.9455e-02, -1.4621e-02, -1.9711e-02,  ..., -2.1926e-02,\n",
+      "         -1.4545e-02, -2.0314e-03]], device='cuda:0'), 'base_model.model.decoder.block.8.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0086, -0.0127,  0.0164,  ..., -0.0085,  0.0160,  0.0239],\n",
+      "        [-0.0253,  0.0110, -0.0184,  ...,  0.0062, -0.0089,  0.0040],\n",
+      "        [ 0.0300, -0.0046, -0.0077,  ...,  0.0109,  0.0151,  0.0012],\n",
+      "        ...,\n",
+      "        [ 0.0069, -0.0402,  0.0182,  ..., -0.0243,  0.0367,  0.0397],\n",
+      "        [ 0.0392, -0.0008,  0.0079,  ..., -0.0269, -0.0036, -0.0002],\n",
+      "        [ 0.0118, -0.0075, -0.0023,  ..., -0.0224, -0.0161,  0.0017]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.8.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0047,  0.0029, -0.0069,  ..., -0.0018, -0.0063, -0.0014],\n",
+      "        [ 0.0082, -0.0077,  0.0052,  ...,  0.0103,  0.0080,  0.0121],\n",
+      "        [-0.0195,  0.0196, -0.0149,  ..., -0.0190, -0.0054, -0.0194],\n",
+      "        ...,\n",
+      "        [-0.0066,  0.0018, -0.0003,  ..., -0.0015, -0.0070,  0.0014],\n",
+      "        [ 0.0169, -0.0169,  0.0163,  ...,  0.0180, -0.0044,  0.0176],\n",
+      "        [ 0.0115, -0.0148,  0.0143,  ...,  0.0089,  0.0165,  0.0110]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.8.layer.1.EncDecAttention.q.lora_A.weight': tensor([[ 0.0088, -0.0198,  0.0352,  ...,  0.0017, -0.0137,  0.0235],\n",
+      "        [ 0.0074,  0.0117,  0.0034,  ..., -0.0414, -0.0173, -0.0157],\n",
+      "        [-0.0269,  0.0180,  0.0023,  ..., -0.0271, -0.0041, -0.0395],\n",
+      "        ...,\n",
+      "        [-0.0052, -0.0061, -0.0074,  ..., -0.0023,  0.0183,  0.0248],\n",
+      "        [-0.0085,  0.0143,  0.0291,  ..., -0.0317,  0.0161, -0.0440],\n",
+      "        [-0.0242,  0.0251,  0.0268,  ..., -0.0222,  0.0009,  0.0223]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.8.layer.1.EncDecAttention.q.lora_B.weight': tensor([[-0.0070, -0.0066, -0.0060,  ...,  0.0057, -0.0022, -0.0063],\n",
+      "        [ 0.0075,  0.0020,  0.0047,  ..., -0.0082, -0.0012,  0.0090],\n",
+      "        [-0.0149, -0.0074, -0.0133,  ...,  0.0114,  0.0035, -0.0127],\n",
+      "        ...,\n",
+      "        [-0.0060, -0.0045, -0.0088,  ...,  0.0093, -0.0050, -0.0093],\n",
+      "        [-0.0044,  0.0054, -0.0040,  ...,  0.0045,  0.0142, -0.0048],\n",
+      "        [ 0.0050,  0.0038,  0.0082,  ..., -0.0098, -0.0018,  0.0068]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.8.layer.1.EncDecAttention.v.lora_A.weight': tensor([[ 0.0147,  0.0348, -0.0170,  ..., -0.0441,  0.0298,  0.0520],\n",
+      "        [-0.0203,  0.0083,  0.0499,  ...,  0.0094,  0.0094, -0.0457],\n",
+      "        [ 0.0123,  0.0063,  0.0027,  ...,  0.0321, -0.0489, -0.0257],\n",
+      "        ...,\n",
+      "        [ 0.0059,  0.0268, -0.0280,  ...,  0.0138, -0.0344, -0.0239],\n",
+      "        [ 0.0025, -0.0047, -0.0095,  ...,  0.0065, -0.0235,  0.0059],\n",
+      "        [-0.0239, -0.0193, -0.0089,  ..., -0.0353, -0.0026,  0.0523]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.8.layer.1.EncDecAttention.v.lora_B.weight': tensor([[-0.0121,  0.0075,  0.0079,  ...,  0.0105,  0.0129, -0.0121],\n",
+      "        [ 0.0100,  0.0012, -0.0099,  ..., -0.0121, -0.0096,  0.0088],\n",
+      "        [ 0.0122, -0.0076, -0.0110,  ..., -0.0120, -0.0115,  0.0121],\n",
+      "        ...,\n",
+      "        [ 0.0093, -0.0080, -0.0108,  ..., -0.0075, -0.0032, -0.0013],\n",
+      "        [-0.0115,  0.0213,  0.0142,  ...,  0.0134,  0.0117, -0.0052],\n",
+      "        [-0.0036,  0.0162,  0.0016,  ...,  0.0018, -0.0023,  0.0021]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.9.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0154,  0.0037,  0.0349,  ..., -0.0158,  0.0038, -0.0385],\n",
+      "        [-0.0073,  0.0473, -0.0104,  ..., -0.0131, -0.0052, -0.0014],\n",
+      "        [ 0.0135, -0.0418,  0.0044,  ...,  0.0084,  0.0183, -0.0136],\n",
+      "        ...,\n",
+      "        [ 0.0269,  0.0013,  0.0371,  ..., -0.0096,  0.0070, -0.0114],\n",
+      "        [-0.0020, -0.0478,  0.0082,  ..., -0.0092,  0.0111, -0.0221],\n",
+      "        [ 0.0107, -0.0073, -0.0224,  ...,  0.0339, -0.0262,  0.0477]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.9.layer.0.SelfAttention.q.lora_B.weight': tensor([[-0.0193,  0.0199, -0.0044,  ..., -0.0231, -0.0193,  0.0192],\n",
+      "        [-0.0032,  0.0015, -0.0103,  ...,  0.0019,  0.0071,  0.0039],\n",
+      "        [-0.0046,  0.0031, -0.0092,  ..., -0.0005,  0.0053,  0.0053],\n",
+      "        ...,\n",
+      "        [-0.0172,  0.0136,  0.0109,  ..., -0.0211, -0.0160,  0.0031],\n",
+      "        [-0.0070,  0.0049, -0.0063,  ..., -0.0058, -0.0043,  0.0051],\n",
+      "        [-0.0118,  0.0162, -0.0009,  ..., -0.0190, -0.0135,  0.0053]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.9.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0438, -0.0172,  0.0363,  ...,  0.0086, -0.0087,  0.0066],\n",
+      "        [ 0.0410, -0.0098,  0.0053,  ..., -0.0250,  0.0264,  0.0233],\n",
+      "        [ 0.0430, -0.0061,  0.0150,  ...,  0.0156,  0.0056, -0.0025],\n",
+      "        ...,\n",
+      "        [ 0.0078, -0.0022,  0.0056,  ..., -0.0328,  0.0030, -0.0042],\n",
+      "        [ 0.0094, -0.0351, -0.0269,  ..., -0.0067,  0.0339,  0.0215],\n",
+      "        [ 0.0293,  0.0221,  0.0090,  ...,  0.0187,  0.0156,  0.0453]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.9.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0210, -0.0218, -0.0219,  ...,  0.0225, -0.0235, -0.0230],\n",
+      "        [-0.0032,  0.0079, -0.0071,  ..., -0.0083,  0.0067,  0.0079],\n",
+      "        [ 0.0085, -0.0037, -0.0048,  ...,  0.0048, -0.0062, -0.0076],\n",
+      "        ...,\n",
+      "        [-0.0054,  0.0032,  0.0100,  ..., -0.0020,  0.0055,  0.0025],\n",
+      "        [-0.0110,  0.0081,  0.0158,  ..., -0.0079,  0.0112,  0.0075],\n",
+      "        [-0.0049,  0.0053,  0.0072,  ..., -0.0049,  0.0068,  0.0042]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.9.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0071, -0.0169, -0.0282,  ...,  0.0228,  0.0082, -0.0134],\n",
+      "        [-0.0093, -0.0039, -0.0467,  ...,  0.0432,  0.0402,  0.0078],\n",
+      "        [ 0.0147,  0.0395,  0.0144,  ..., -0.0338, -0.0109,  0.0084],\n",
+      "        ...,\n",
+      "        [ 0.0166, -0.0008,  0.0265,  ..., -0.0120, -0.0385, -0.0455],\n",
+      "        [ 0.0310, -0.0093,  0.0134,  ..., -0.0040, -0.0034, -0.0137],\n",
+      "        [-0.0219, -0.0405, -0.0077,  ...,  0.0289, -0.0082, -0.0013]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.9.layer.1.EncDecAttention.q.lora_B.weight': tensor([[ 0.0050, -0.0137,  0.0062,  ...,  0.0020, -0.0027, -0.0102],\n",
+      "        [-0.0208,  0.0056, -0.0130,  ..., -0.0062,  0.0127,  0.0047],\n",
+      "        [-0.0034, -0.0049,  0.0036,  ...,  0.0025,  0.0036,  0.0013],\n",
+      "        ...,\n",
+      "        [ 0.0123, -0.0268,  0.0155,  ...,  0.0150, -0.0170, -0.0189],\n",
+      "        [-0.0296,  0.0265, -0.0202,  ..., -0.0150,  0.0275,  0.0163],\n",
+      "        [-0.0052,  0.0020, -0.0008,  ..., -0.0013,  0.0061,  0.0013]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.9.layer.1.EncDecAttention.v.lora_A.weight': tensor([[ 0.0483,  0.0039,  0.0312,  ..., -0.0194,  0.0217,  0.0157],\n",
+      "        [-0.0274,  0.0153,  0.0293,  ...,  0.0371, -0.0078, -0.0204],\n",
+      "        [-0.0226,  0.0241, -0.0100,  ...,  0.0308, -0.0354, -0.0683],\n",
+      "        ...,\n",
+      "        [ 0.0043,  0.0249, -0.0220,  ..., -0.0079,  0.0153,  0.0055],\n",
+      "        [ 0.0618,  0.0091,  0.0140,  ..., -0.0128,  0.0089,  0.0166],\n",
+      "        [-0.0502, -0.0238,  0.0167,  ...,  0.0662,  0.0163, -0.0201]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.9.layer.1.EncDecAttention.v.lora_B.weight': tensor([[-0.0136,  0.0160,  0.0156,  ..., -0.0165, -0.0171,  0.0161],\n",
+      "        [ 0.0025, -0.0026, -0.0052,  ...,  0.0055,  0.0045, -0.0020],\n",
+      "        [-0.0057,  0.0066,  0.0048,  ..., -0.0033, -0.0056,  0.0068],\n",
+      "        ...,\n",
+      "        [ 0.0240, -0.0242, -0.0254,  ...,  0.0231,  0.0241, -0.0242],\n",
+      "        [ 0.0213, -0.0170, -0.0214,  ...,  0.0198,  0.0180, -0.0202],\n",
+      "        [-0.0150,  0.0218,  0.0179,  ..., -0.0128, -0.0194,  0.0167]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.10.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0025, -0.0290,  0.0255,  ..., -0.0181,  0.0138,  0.0025],\n",
+      "        [-0.0440, -0.0035,  0.0296,  ...,  0.0243,  0.0088, -0.0053],\n",
+      "        [-0.0272,  0.0054, -0.0146,  ...,  0.0257,  0.0243,  0.0020],\n",
+      "        ...,\n",
+      "        [ 0.0010,  0.0344,  0.0134,  ...,  0.0117, -0.0007, -0.0075],\n",
+      "        [ 0.0423, -0.0215,  0.0008,  ..., -0.0328, -0.0238, -0.0170],\n",
+      "        [-0.0181,  0.0097,  0.0098,  ...,  0.0165, -0.0156, -0.0103]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.10.layer.0.SelfAttention.q.lora_B.weight': tensor([[-0.0063,  0.0127,  0.0105,  ...,  0.0136,  0.0137,  0.0123],\n",
+      "        [ 0.0126, -0.0181, -0.0162,  ..., -0.0186, -0.0186, -0.0181],\n",
+      "        [-0.0070,  0.0156,  0.0164,  ...,  0.0148,  0.0144,  0.0143],\n",
+      "        ...,\n",
+      "        [-0.0039,  0.0014,  0.0003,  ...,  0.0025, -0.0060,  0.0027],\n",
+      "        [ 0.0025,  0.0164,  0.0130,  ...,  0.0138, -0.0010,  0.0125],\n",
+      "        [ 0.0102,  0.0055,  0.0062,  ..., -0.0012, -0.0060,  0.0018]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.10.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0024, -0.0390, -0.0036,  ...,  0.0109, -0.0263,  0.0339],\n",
+      "        [ 0.0104, -0.0131,  0.0153,  ..., -0.0151,  0.0249,  0.0227],\n",
+      "        [-0.0228,  0.0018,  0.0311,  ..., -0.0179,  0.0120, -0.0137],\n",
+      "        ...,\n",
+      "        [ 0.0398, -0.0251,  0.0051,  ...,  0.0207,  0.0257,  0.0066],\n",
+      "        [ 0.0026,  0.0058, -0.0247,  ...,  0.0021,  0.0233, -0.0101],\n",
+      "        [ 0.0265,  0.0026,  0.0404,  ...,  0.0282, -0.0161,  0.0450]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.10.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0041,  0.0035, -0.0041,  ...,  0.0040,  0.0048,  0.0029],\n",
+      "        [-0.0153, -0.0180,  0.0142,  ..., -0.0159, -0.0138, -0.0151],\n",
+      "        [ 0.0144,  0.0131, -0.0153,  ...,  0.0147,  0.0151,  0.0122],\n",
+      "        ...,\n",
+      "        [-0.0130, -0.0103,  0.0154,  ..., -0.0113, -0.0158, -0.0153],\n",
+      "        [-0.0170, -0.0153,  0.0144,  ..., -0.0125, -0.0155, -0.0185],\n",
+      "        [-0.0158, -0.0154,  0.0174,  ..., -0.0161, -0.0178, -0.0153]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.10.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0304, -0.0075, -0.0230,  ...,  0.0063, -0.0018,  0.0278],\n",
+      "        [-0.0076, -0.0079,  0.0015,  ..., -0.0316,  0.0023,  0.0243],\n",
+      "        [ 0.0084,  0.0124,  0.0012,  ...,  0.0109, -0.0317,  0.0220],\n",
+      "        ...,\n",
+      "        [-0.0283,  0.0414, -0.0127,  ..., -0.0159,  0.0182,  0.0143],\n",
+      "        [-0.0043, -0.0115, -0.0229,  ...,  0.0092, -0.0283,  0.0275],\n",
+      "        [-0.0176,  0.0361, -0.0088,  ..., -0.0113, -0.0109,  0.0135]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.10.layer.1.EncDecAttention.q.lora_B.weight': tensor([[ 8.9059e-03,  5.9997e-03,  3.2296e-03,  ...,  4.1912e-03,\n",
+      "         -7.3702e-04,  3.2630e-03],\n",
+      "        [ 1.4178e-02,  1.5074e-02,  8.4570e-03,  ...,  3.8711e-03,\n",
+      "          1.0596e-02,  6.3549e-03],\n",
+      "        [-1.2284e-02, -9.7453e-03,  5.3009e-04,  ...,  8.9120e-03,\n",
+      "          5.1250e-04, -5.0001e-03],\n",
+      "        ...,\n",
+      "        [ 4.1886e-03,  8.0591e-05, -2.7274e-04,  ...,  8.2174e-03,\n",
+      "         -2.2814e-03,  4.5755e-04],\n",
+      "        [ 1.7286e-02,  1.7134e-02,  1.1415e-02,  ...,  2.6513e-02,\n",
+      "          2.5976e-03,  1.1181e-02],\n",
+      "        [ 1.2990e-02,  1.4693e-02,  1.0798e-02,  ...,  1.6177e-02,\n",
+      "          5.8300e-03,  5.1070e-03]], device='cuda:0'), 'base_model.model.decoder.block.10.layer.1.EncDecAttention.v.lora_A.weight': tensor([[ 0.0020, -0.0097, -0.0073,  ..., -0.0110,  0.0074,  0.0097],\n",
+      "        [-0.0227,  0.0208, -0.0257,  ..., -0.0031, -0.0147, -0.0224],\n",
+      "        [-0.0020,  0.0450,  0.0228,  ..., -0.0019, -0.0001,  0.0342],\n",
+      "        ...,\n",
+      "        [-0.0067, -0.0084,  0.0036,  ...,  0.0493, -0.0207, -0.0334],\n",
+      "        [ 0.0320, -0.0389, -0.0380,  ...,  0.0327, -0.0444,  0.0107],\n",
+      "        [-0.0115, -0.0011,  0.0172,  ...,  0.0082, -0.0172, -0.0654]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.10.layer.1.EncDecAttention.v.lora_B.weight': tensor([[ 0.0213,  0.0044,  0.0228,  ..., -0.0223, -0.0208, -0.0105],\n",
+      "        [-0.0181,  0.0040, -0.0209,  ...,  0.0195,  0.0192,  0.0131],\n",
+      "        [ 0.0149,  0.0117,  0.0133,  ..., -0.0163, -0.0158, -0.0026],\n",
+      "        ...,\n",
+      "        [ 0.0104,  0.0159,  0.0134,  ..., -0.0119, -0.0154, -0.0064],\n",
+      "        [-0.0195, -0.0169, -0.0204,  ...,  0.0204,  0.0179,  0.0020],\n",
+      "        [-0.0017,  0.0049, -0.0006,  ..., -0.0002,  0.0053,  0.0112]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.11.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0113, -0.0173, -0.0027,  ...,  0.0034,  0.0462, -0.0084],\n",
+      "        [-0.0225, -0.0064, -0.0297,  ..., -0.0134,  0.0226,  0.0151],\n",
+      "        [ 0.0253, -0.0319,  0.0379,  ...,  0.0115,  0.0016, -0.0272],\n",
+      "        ...,\n",
+      "        [ 0.0177,  0.0429, -0.0232,  ...,  0.0427,  0.0081,  0.0265],\n",
+      "        [-0.0095, -0.0153, -0.0117,  ..., -0.0021, -0.0487, -0.0410],\n",
+      "        [ 0.0021,  0.0395,  0.0165,  ...,  0.0390,  0.0358,  0.0250]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.11.layer.0.SelfAttention.q.lora_B.weight': tensor([[-0.0042, -0.0051,  0.0075,  ..., -0.0086,  0.0048, -0.0102],\n",
+      "        [-0.0015, -0.0072,  0.0053,  ..., -0.0051,  0.0027, -0.0075],\n",
+      "        [ 0.0172,  0.0165, -0.0048,  ...,  0.0060, -0.0144,  0.0042],\n",
+      "        ...,\n",
+      "        [-0.0212, -0.0192,  0.0205,  ..., -0.0172,  0.0160, -0.0143],\n",
+      "        [ 0.0145,  0.0102, -0.0123,  ...,  0.0128, -0.0037,  0.0073],\n",
+      "        [-0.0072, -0.0060,  0.0064,  ..., -0.0072,  0.0052, -0.0087]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.11.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0115, -0.0157,  0.0200,  ..., -0.0234,  0.0127,  0.0111],\n",
+      "        [ 0.0194,  0.0063, -0.0089,  ..., -0.0268, -0.0184, -0.0171],\n",
+      "        [ 0.0094, -0.0173,  0.0352,  ..., -0.0146, -0.0239,  0.0131],\n",
+      "        ...,\n",
+      "        [ 0.0015, -0.0065, -0.0278,  ..., -0.0262,  0.0022, -0.0093],\n",
+      "        [ 0.0198, -0.0332, -0.0144,  ...,  0.0026,  0.0404,  0.0179],\n",
+      "        [ 0.0105, -0.0364,  0.0262,  ..., -0.0171,  0.0403,  0.0121]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.11.layer.0.SelfAttention.v.lora_B.weight': tensor([[-9.5707e-03, -1.2350e-02,  4.2512e-03,  ..., -6.5697e-03,\n",
+      "          1.0621e-02,  8.1075e-03],\n",
+      "        [-1.3202e-03,  1.8115e-02, -1.6808e-02,  ...,  1.9649e-02,\n",
+      "         -1.0918e-02, -1.9214e-02],\n",
+      "        [ 1.8847e-02, -1.7127e-02,  1.7074e-02,  ..., -1.9088e-02,\n",
+      "          8.5360e-03,  1.7337e-02],\n",
+      "        ...,\n",
+      "        [-2.9112e-03, -2.2279e-02,  2.3047e-02,  ..., -2.6989e-02,\n",
+      "          2.1065e-02,  2.5149e-02],\n",
+      "        [ 1.2249e-02, -5.4296e-04,  1.3645e-02,  ..., -7.0064e-04,\n",
+      "         -1.2102e-03, -8.9333e-05],\n",
+      "        [-3.2261e-03,  1.1069e-02, -3.9940e-03,  ...,  1.3339e-02,\n",
+      "         -1.5026e-02, -1.2184e-02]], device='cuda:0'), 'base_model.model.decoder.block.11.layer.1.EncDecAttention.q.lora_A.weight': tensor([[ 0.0151,  0.0504,  0.0151,  ..., -0.0405, -0.0372,  0.0045],\n",
+      "        [ 0.0058, -0.0132, -0.0114,  ...,  0.0315, -0.0070, -0.0143],\n",
+      "        [ 0.0007, -0.0070, -0.0318,  ..., -0.0215,  0.0071,  0.0102],\n",
+      "        ...,\n",
+      "        [-0.0004,  0.0088,  0.0068,  ...,  0.0079, -0.0094,  0.0206],\n",
+      "        [-0.0109,  0.0159, -0.0031,  ..., -0.0096, -0.0329, -0.0108],\n",
+      "        [-0.0080,  0.0129,  0.0347,  ..., -0.0115,  0.0031,  0.0019]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.11.layer.1.EncDecAttention.q.lora_B.weight': tensor([[-0.0131,  0.0158,  0.0049,  ...,  0.0131, -0.0112, -0.0145],\n",
+      "        [ 0.0174, -0.0180, -0.0168,  ..., -0.0093,  0.0182,  0.0186],\n",
+      "        [-0.0105,  0.0116,  0.0139,  ..., -0.0073, -0.0136, -0.0119],\n",
+      "        ...,\n",
+      "        [-0.0058,  0.0093,  0.0242,  ..., -0.0101, -0.0032,  0.0060],\n",
+      "        [ 0.0191, -0.0159, -0.0204,  ...,  0.0079,  0.0200,  0.0092],\n",
+      "        [ 0.0009, -0.0016,  0.0067,  ...,  0.0027, -0.0018,  0.0032]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.11.layer.1.EncDecAttention.v.lora_A.weight': tensor([[ 0.0233,  0.0394, -0.0297,  ...,  0.0140, -0.0410, -0.0096],\n",
+      "        [-0.0128, -0.0103,  0.0050,  ..., -0.0203,  0.0273, -0.0021],\n",
+      "        [ 0.0241,  0.0009, -0.0340,  ..., -0.0166, -0.0095,  0.0592],\n",
+      "        ...,\n",
+      "        [ 0.0344,  0.0115, -0.0103,  ..., -0.0338,  0.0280,  0.0553],\n",
+      "        [-0.0379, -0.0005,  0.0447,  ..., -0.0204,  0.0497, -0.0573],\n",
+      "        [ 0.0310, -0.0138, -0.0048,  ...,  0.0067, -0.0408,  0.0392]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.11.layer.1.EncDecAttention.v.lora_B.weight': tensor([[ 2.2074e-02, -2.3050e-02, -2.0557e-02,  ..., -3.0574e-02,\n",
+      "          6.5095e-04,  7.4577e-03],\n",
+      "        [ 1.9601e-02, -1.9351e-02, -1.3357e-02,  ..., -1.9333e-02,\n",
+      "         -5.8716e-03,  1.7343e-02],\n",
+      "        [-2.0707e-02,  2.1192e-02,  1.7973e-02,  ...,  2.2283e-02,\n",
+      "          7.3185e-03, -1.4129e-02],\n",
+      "        ...,\n",
+      "        [-1.2650e-02,  5.2220e-03,  2.7920e-04,  ...,  1.1320e-03,\n",
+      "          2.8257e-03, -8.0793e-03],\n",
+      "        [-9.3410e-03, -1.6412e-03, -1.1426e-02,  ..., -1.1142e-02,\n",
+      "          5.2212e-03, -1.3261e-03],\n",
+      "        [-1.8118e-02,  1.8327e-02, -1.9824e-05,  ...,  2.1927e-02,\n",
+      "          8.2317e-03, -1.1294e-02]], device='cuda:0'), 'base_model.model.decoder.block.12.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0103,  0.0235, -0.0445,  ..., -0.0011,  0.0211,  0.0054],\n",
+      "        [ 0.0231,  0.0272, -0.0111,  ..., -0.0416, -0.0326, -0.0198],\n",
+      "        [-0.0178, -0.0060,  0.0432,  ..., -0.0134,  0.0074,  0.0074],\n",
+      "        ...,\n",
+      "        [ 0.0253,  0.0366,  0.0049,  ..., -0.0379, -0.0321,  0.0107],\n",
+      "        [-0.0126, -0.0017, -0.0359,  ...,  0.0121,  0.0077,  0.0259],\n",
+      "        [ 0.0278,  0.0117, -0.0050,  ...,  0.0169,  0.0003,  0.0109]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.12.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0189, -0.0169, -0.0191,  ..., -0.0217,  0.0187, -0.0174],\n",
+      "        [-0.0155,  0.0188,  0.0232,  ...,  0.0169, -0.0256,  0.0255],\n",
+      "        [-0.0193,  0.0153,  0.0159,  ...,  0.0276, -0.0155,  0.0142],\n",
+      "        ...,\n",
+      "        [-0.0200,  0.0177,  0.0180,  ...,  0.0193, -0.0177,  0.0189],\n",
+      "        [-0.0146,  0.0139,  0.0137,  ...,  0.0132, -0.0135,  0.0142],\n",
+      "        [ 0.0102, -0.0118, -0.0130,  ..., -0.0131,  0.0132, -0.0101]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.12.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 2.4031e-03,  2.1623e-02,  9.9847e-03,  ...,  1.5630e-02,\n",
+      "         -8.4626e-03, -1.9406e-02],\n",
+      "        [ 6.6478e-03,  2.4983e-02, -1.5062e-02,  ...,  1.9716e-02,\n",
+      "         -1.7178e-02,  1.3483e-03],\n",
+      "        [ 9.8836e-03,  6.9536e-03,  2.2963e-02,  ...,  3.3474e-02,\n",
+      "          2.1320e-02,  3.9549e-02],\n",
+      "        ...,\n",
+      "        [-2.3796e-02, -7.5400e-06,  3.0448e-02,  ...,  4.2281e-02,\n",
+      "          2.4435e-02,  3.4439e-03],\n",
+      "        [-8.6579e-03,  2.4908e-02, -1.4695e-02,  ..., -1.7937e-02,\n",
+      "         -2.9916e-02, -7.0838e-03],\n",
+      "        [-1.8854e-02, -4.6361e-02,  4.9669e-02,  ..., -1.4553e-04,\n",
+      "          2.3075e-02,  1.9088e-03]], device='cuda:0'), 'base_model.model.decoder.block.12.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0074, -0.0017,  0.0112,  ...,  0.0109, -0.0103,  0.0119],\n",
+      "        [-0.0149, -0.0008,  0.0157,  ...,  0.0172, -0.0093,  0.0181],\n",
+      "        [ 0.0007,  0.0051, -0.0025,  ..., -0.0014,  0.0031, -0.0015],\n",
+      "        ...,\n",
+      "        [-0.0013, -0.0033, -0.0026,  ..., -0.0015, -0.0034, -0.0038],\n",
+      "        [-0.0056, -0.0132,  0.0050,  ...,  0.0036, -0.0130, -0.0087],\n",
+      "        [-0.0098, -0.0120,  0.0107,  ...,  0.0110, -0.0117,  0.0023]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.12.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0046,  0.0211,  0.0200,  ..., -0.0087,  0.0239,  0.0120],\n",
+      "        [-0.0436, -0.0509, -0.0096,  ..., -0.0034,  0.0071, -0.0181],\n",
+      "        [-0.0019, -0.0168, -0.0081,  ..., -0.0142, -0.0002,  0.0396],\n",
+      "        ...,\n",
+      "        [ 0.0069, -0.0059,  0.0063,  ..., -0.0305,  0.0049,  0.0234],\n",
+      "        [-0.0008, -0.0510, -0.0299,  ...,  0.0231,  0.0161,  0.0187],\n",
+      "        [ 0.0062, -0.0050, -0.0067,  ...,  0.0040, -0.0048, -0.0111]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.12.layer.1.EncDecAttention.q.lora_B.weight': tensor([[ 0.0257, -0.0312, -0.0249,  ..., -0.0225, -0.0249,  0.0227],\n",
+      "        [ 0.0275, -0.0350, -0.0293,  ..., -0.0257, -0.0283,  0.0259],\n",
+      "        [ 0.0032, -0.0025,  0.0020,  ..., -0.0031, -0.0025,  0.0046],\n",
+      "        ...,\n",
+      "        [-0.0013,  0.0057,  0.0068,  ...,  0.0017,  0.0016, -0.0019],\n",
+      "        [ 0.0046, -0.0059, -0.0001,  ..., -0.0073, -0.0064,  0.0085],\n",
+      "        [-0.0058,  0.0037,  0.0097,  ...,  0.0026,  0.0047, -0.0013]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.12.layer.1.EncDecAttention.v.lora_A.weight': tensor([[ 0.0103, -0.0327, -0.0249,  ...,  0.0200, -0.0102,  0.0046],\n",
+      "        [ 0.0150,  0.0122, -0.0172,  ...,  0.0032,  0.0125, -0.0109],\n",
+      "        [-0.0103,  0.0146,  0.0295,  ..., -0.0434, -0.0173, -0.0248],\n",
+      "        ...,\n",
+      "        [ 0.0349, -0.0145,  0.0197,  ...,  0.0345, -0.0413,  0.0070],\n",
+      "        [-0.0015, -0.0148,  0.0261,  ..., -0.0024, -0.0134, -0.0016],\n",
+      "        [ 0.0220,  0.0160,  0.0229,  ..., -0.0040, -0.0127,  0.0081]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.12.layer.1.EncDecAttention.v.lora_B.weight': tensor([[ 0.0108,  0.0126, -0.0129,  ...,  0.0061, -0.0068,  0.0116],\n",
+      "        [-0.0031, -0.0131,  0.0072,  ..., -0.0232,  0.0122, -0.0168],\n",
+      "        [-0.0215, -0.0097,  0.0125,  ..., -0.0147,  0.0112, -0.0142],\n",
+      "        ...,\n",
+      "        [-0.0178, -0.0196,  0.0217,  ..., -0.0212,  0.0185, -0.0224],\n",
+      "        [ 0.0107, -0.0059,  0.0002,  ..., -0.0074,  0.0075, -0.0123],\n",
+      "        [ 0.0042, -0.0160, -0.0046,  ..., -0.0154,  0.0115, -0.0159]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.13.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0219,  0.0094, -0.0227,  ..., -0.0118, -0.0069, -0.0019],\n",
+      "        [ 0.0170,  0.0005,  0.0317,  ..., -0.0442, -0.0148, -0.0178],\n",
+      "        [-0.0057, -0.0064,  0.0218,  ..., -0.0492, -0.0141, -0.0176],\n",
+      "        ...,\n",
+      "        [ 0.0341, -0.0037,  0.0090,  ...,  0.0453,  0.0237,  0.0424],\n",
+      "        [ 0.0107, -0.0036,  0.0088,  ..., -0.0421, -0.0325, -0.0168],\n",
+      "        [ 0.0242, -0.0061, -0.0386,  ..., -0.0047, -0.0117, -0.0059]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.13.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0086, -0.0127, -0.0096,  ...,  0.0039, -0.0092,  0.0090],\n",
+      "        [ 0.0119, -0.0163, -0.0191,  ...,  0.0047, -0.0144,  0.0121],\n",
+      "        [-0.0091,  0.0139,  0.0143,  ..., -0.0050,  0.0106, -0.0094],\n",
+      "        ...,\n",
+      "        [ 0.0141, -0.0123, -0.0156,  ...,  0.0110, -0.0122,  0.0114],\n",
+      "        [-0.0024,  0.0034,  0.0031,  ..., -0.0013,  0.0014, -0.0071],\n",
+      "        [ 0.0017, -0.0023, -0.0021,  ..., -0.0023,  0.0005,  0.0037]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.13.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0047,  0.0136, -0.0285,  ..., -0.0433,  0.0030,  0.0144],\n",
+      "        [ 0.0114,  0.0227, -0.0303,  ...,  0.0104,  0.0045,  0.0162],\n",
+      "        [ 0.0358,  0.0059,  0.0081,  ..., -0.0286, -0.0211, -0.0286],\n",
+      "        ...,\n",
+      "        [ 0.0013, -0.0096,  0.0008,  ...,  0.0270, -0.0111,  0.0180],\n",
+      "        [ 0.0205,  0.0279, -0.0179,  ..., -0.0381, -0.0168, -0.0347],\n",
+      "        [ 0.0088, -0.0382,  0.0431,  ...,  0.0359, -0.0130,  0.0356]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.13.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0092, -0.0092, -0.0036,  ...,  0.0120, -0.0077,  0.0126],\n",
+      "        [ 0.0039,  0.0076,  0.0078,  ...,  0.0046,  0.0068,  0.0035],\n",
+      "        [-0.0306, -0.0290, -0.0167,  ...,  0.0265, -0.0283,  0.0298],\n",
+      "        ...,\n",
+      "        [ 0.0154,  0.0149,  0.0118,  ..., -0.0089,  0.0155, -0.0134],\n",
+      "        [-0.0122, -0.0120, -0.0107,  ...,  0.0109, -0.0126,  0.0120],\n",
+      "        [-0.0047, -0.0033,  0.0010,  ...,  0.0054,  0.0011,  0.0038]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.13.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0253, -0.0134, -0.0078,  ..., -0.0472, -0.0238, -0.0436],\n",
+      "        [ 0.0033, -0.0051,  0.0008,  ...,  0.0166, -0.0053,  0.0373],\n",
+      "        [-0.0360, -0.0351, -0.0227,  ...,  0.0108,  0.0331, -0.0146],\n",
+      "        ...,\n",
+      "        [-0.0030,  0.0173, -0.0300,  ...,  0.0246,  0.0367,  0.0359],\n",
+      "        [ 0.0126,  0.0130,  0.0141,  ...,  0.0443, -0.0219, -0.0240],\n",
+      "        [ 0.0219,  0.0208, -0.0084,  ...,  0.0185,  0.0198,  0.0202]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.13.layer.1.EncDecAttention.q.lora_B.weight': tensor([[-0.0225,  0.0218,  0.0237,  ...,  0.0197,  0.0006,  0.0226],\n",
+      "        [-0.0211,  0.0232,  0.0246,  ...,  0.0203,  0.0224,  0.0244],\n",
+      "        [-0.0058, -0.0001,  0.0041,  ...,  0.0010,  0.0181,  0.0170],\n",
+      "        ...,\n",
+      "        [ 0.0031,  0.0004,  0.0067,  ..., -0.0013,  0.0070,  0.0063],\n",
+      "        [ 0.0117, -0.0107, -0.0180,  ..., -0.0149, -0.0202, -0.0204],\n",
+      "        [-0.0079,  0.0058,  0.0041,  ...,  0.0028,  0.0070,  0.0060]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.13.layer.1.EncDecAttention.v.lora_A.weight': tensor([[-0.0245, -0.0014,  0.0118,  ...,  0.0183,  0.0520,  0.0118],\n",
+      "        [-0.0252,  0.0178,  0.0142,  ...,  0.0115,  0.0291,  0.0260],\n",
+      "        [-0.0256,  0.0298,  0.0250,  ..., -0.0366,  0.0294,  0.0026],\n",
+      "        ...,\n",
+      "        [ 0.0112,  0.0158, -0.0059,  ...,  0.0497, -0.0213, -0.0500],\n",
+      "        [-0.0259,  0.0207, -0.0143,  ..., -0.0326,  0.0212,  0.0302],\n",
+      "        [-0.0123,  0.0010, -0.0071,  ...,  0.0158, -0.0003, -0.0120]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.13.layer.1.EncDecAttention.v.lora_B.weight': tensor([[ 0.0100,  0.0163,  0.0209,  ..., -0.0165,  0.0191,  0.0183],\n",
+      "        [-0.0046, -0.0019, -0.0039,  ...,  0.0066, -0.0077, -0.0095],\n",
+      "        [ 0.0071,  0.0146,  0.0189,  ..., -0.0203,  0.0199,  0.0203],\n",
+      "        ...,\n",
+      "        [ 0.0047,  0.0073,  0.0046,  ..., -0.0153,  0.0060,  0.0022],\n",
+      "        [-0.0199, -0.0175, -0.0212,  ...,  0.0074, -0.0213, -0.0187],\n",
+      "        [-0.0100, -0.0072, -0.0100,  ...,  0.0194, -0.0130, -0.0081]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.14.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0374, -0.0102, -0.0033,  ...,  0.0444, -0.0081, -0.0211],\n",
+      "        [ 0.0251, -0.0035, -0.0038,  ..., -0.0019, -0.0239,  0.0165],\n",
+      "        [ 0.0151,  0.0200, -0.0244,  ...,  0.0094, -0.0323, -0.0069],\n",
+      "        ...,\n",
+      "        [ 0.0082,  0.0296, -0.0086,  ...,  0.0067, -0.0011,  0.0200],\n",
+      "        [ 0.0051, -0.0162,  0.0073,  ..., -0.0121, -0.0001, -0.0147],\n",
+      "        [-0.0156,  0.0238,  0.0216,  ..., -0.0370, -0.0026, -0.0052]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.14.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0076, -0.0004, -0.0016,  ...,  0.0082,  0.0001,  0.0090],\n",
+      "        [-0.0125, -0.0170, -0.0166,  ..., -0.0114,  0.0149, -0.0104],\n",
+      "        [ 0.0138,  0.0154,  0.0162,  ...,  0.0136, -0.0150,  0.0128],\n",
+      "        ...,\n",
+      "        [ 0.0039,  0.0059,  0.0069,  ...,  0.0038, -0.0016,  0.0073],\n",
+      "        [ 0.0080, -0.0075, -0.0050,  ..., -0.0102,  0.0071, -0.0042],\n",
+      "        [ 0.0025,  0.0050,  0.0025,  ...,  0.0080, -0.0047,  0.0008]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.14.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0017,  0.0335,  0.0202,  ..., -0.0248,  0.0198, -0.0257],\n",
+      "        [ 0.0159, -0.0008,  0.0156,  ...,  0.0062,  0.0148,  0.0142],\n",
+      "        [ 0.0241,  0.0451, -0.0204,  ..., -0.0211, -0.0325, -0.0450],\n",
+      "        ...,\n",
+      "        [ 0.0282,  0.0004, -0.0164,  ..., -0.0083, -0.0320,  0.0056],\n",
+      "        [-0.0371, -0.0190, -0.0167,  ...,  0.0097,  0.0287,  0.0375],\n",
+      "        [ 0.0385,  0.0255, -0.0023,  ..., -0.0265, -0.0068, -0.0086]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.14.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0089, -0.0147, -0.0117,  ..., -0.0144,  0.0151, -0.0129],\n",
+      "        [ 0.0181,  0.0294,  0.0288,  ...,  0.0267, -0.0299,  0.0301],\n",
+      "        [-0.0206, -0.0198, -0.0207,  ..., -0.0200,  0.0170, -0.0197],\n",
+      "        ...,\n",
+      "        [-0.0094,  0.0028,  0.0006,  ...,  0.0009, -0.0027, -0.0012],\n",
+      "        [-0.0130, -0.0167, -0.0158,  ..., -0.0158,  0.0143, -0.0159],\n",
+      "        [-0.0003,  0.0117,  0.0091,  ...,  0.0089, -0.0094,  0.0117]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.14.layer.1.EncDecAttention.q.lora_A.weight': tensor([[ 0.0170, -0.0165, -0.0077,  ...,  0.0260,  0.0332, -0.0158],\n",
+      "        [ 0.0055, -0.0062,  0.0312,  ...,  0.0198,  0.0379, -0.0031],\n",
+      "        [-0.0055, -0.0003, -0.0136,  ...,  0.0274,  0.0339,  0.0083],\n",
+      "        ...,\n",
+      "        [-0.0270, -0.0042,  0.0200,  ..., -0.0558, -0.0361, -0.0222],\n",
+      "        [-0.0186,  0.0363, -0.0226,  ..., -0.0413, -0.0183,  0.0342],\n",
+      "        [ 0.0233, -0.0100, -0.0039,  ...,  0.0253,  0.0414, -0.0166]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.14.layer.1.EncDecAttention.q.lora_B.weight': tensor([[ 0.0069,  0.0116,  0.0013,  ..., -0.0092, -0.0087,  0.0109],\n",
+      "        [-0.0145, -0.0106, -0.0112,  ...,  0.0108,  0.0059, -0.0067],\n",
+      "        [ 0.0009,  0.0047,  0.0026,  ..., -0.0067, -0.0045,  0.0030],\n",
+      "        ...,\n",
+      "        [-0.0020,  0.0029,  0.0040,  ..., -0.0054, -0.0019,  0.0033],\n",
+      "        [ 0.0104,  0.0127,  0.0050,  ..., -0.0121, -0.0120,  0.0125],\n",
+      "        [ 0.0132,  0.0152,  0.0163,  ..., -0.0161, -0.0142,  0.0154]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.14.layer.1.EncDecAttention.v.lora_A.weight': tensor([[ 0.0538, -0.0210, -0.0223,  ...,  0.0296,  0.0070,  0.0165],\n",
+      "        [ 0.0071,  0.0046, -0.0163,  ..., -0.0276, -0.0093,  0.0149],\n",
+      "        [ 0.0524,  0.0300, -0.0185,  ...,  0.0410, -0.0501, -0.0439],\n",
+      "        ...,\n",
+      "        [ 0.0027, -0.0458,  0.0228,  ..., -0.0456,  0.0372, -0.0044],\n",
+      "        [-0.0516, -0.0328,  0.0084,  ..., -0.0020,  0.0101, -0.0217],\n",
+      "        [-0.0488, -0.0184, -0.0100,  ..., -0.0236, -0.0005,  0.0046]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.14.layer.1.EncDecAttention.v.lora_B.weight': tensor([[ 0.0304, -0.0266,  0.0291,  ..., -0.0293, -0.0279, -0.0280],\n",
+      "        [ 0.0024, -0.0023,  0.0030,  ..., -0.0006,  0.0019, -0.0014],\n",
+      "        [-0.0144,  0.0198, -0.0207,  ...,  0.0224,  0.0214,  0.0192],\n",
+      "        ...,\n",
+      "        [ 0.0128, -0.0158,  0.0149,  ..., -0.0154, -0.0151, -0.0157],\n",
+      "        [ 0.0039, -0.0054,  0.0027,  ..., -0.0020, -0.0050, -0.0052],\n",
+      "        [ 0.0104, -0.0127,  0.0108,  ..., -0.0100, -0.0103, -0.0124]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.15.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0160, -0.0093,  0.0195,  ..., -0.0205, -0.0025, -0.0271],\n",
+      "        [ 0.0169, -0.0051,  0.0047,  ...,  0.0031,  0.0108,  0.0349],\n",
+      "        [-0.0346,  0.0079,  0.0012,  ..., -0.0336, -0.0087, -0.0240],\n",
+      "        ...,\n",
+      "        [ 0.0279, -0.0023, -0.0044,  ...,  0.0205, -0.0074,  0.0415],\n",
+      "        [ 0.0306, -0.0036, -0.0323,  ..., -0.0290,  0.0356,  0.0230],\n",
+      "        [-0.0206,  0.0021,  0.0137,  ..., -0.0037,  0.0064, -0.0108]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.15.layer.0.SelfAttention.q.lora_B.weight': tensor([[-0.0021,  0.0034, -0.0017,  ...,  0.0019,  0.0005, -0.0005],\n",
+      "        [ 0.0027,  0.0014,  0.0029,  ..., -0.0032, -0.0004,  0.0007],\n",
+      "        [ 0.0128, -0.0122,  0.0117,  ..., -0.0105, -0.0078,  0.0099],\n",
+      "        ...,\n",
+      "        [ 0.0084, -0.0047,  0.0049,  ..., -0.0058, -0.0098,  0.0095],\n",
+      "        [ 0.0076, -0.0041,  0.0042,  ..., -0.0087, -0.0084,  0.0071],\n",
+      "        [-0.0067, -0.0047, -0.0006,  ...,  0.0078,  0.0089, -0.0042]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.15.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0054, -0.0268,  0.0230,  ...,  0.0257, -0.0094, -0.0198],\n",
+      "        [-0.0194,  0.0296,  0.0148,  ..., -0.0157, -0.0334, -0.0433],\n",
+      "        [-0.0143,  0.0536,  0.0044,  ..., -0.0396, -0.0147, -0.0384],\n",
+      "        ...,\n",
+      "        [-0.0108,  0.0081,  0.0216,  ..., -0.0186, -0.0146,  0.0007],\n",
+      "        [-0.0073, -0.0274, -0.0106,  ..., -0.0029, -0.0138,  0.0210],\n",
+      "        [-0.0089, -0.0516,  0.0305,  ...,  0.0103,  0.0208, -0.0060]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.15.layer.0.SelfAttention.v.lora_B.weight': tensor([[-7.7887e-03, -4.4329e-03, -1.0258e-02,  ..., -1.1257e-02,\n",
+      "          9.5218e-03,  1.0012e-02],\n",
+      "        [-1.3725e-02, -1.3488e-02, -2.1645e-03,  ..., -7.9241e-04,\n",
+      "          2.8600e-03,  2.6073e-03],\n",
+      "        [-1.9941e-02, -5.7291e-03, -3.8434e-03,  ..., -4.3191e-03,\n",
+      "          7.2316e-03,  3.0308e-03],\n",
+      "        ...,\n",
+      "        [-2.7335e-03, -2.2611e-02, -1.1037e-02,  ..., -9.8621e-03,\n",
+      "          1.4747e-02,  1.1001e-02],\n",
+      "        [-1.6099e-02, -1.4242e-03, -2.9456e-03,  ...,  6.2675e-04,\n",
+      "          5.8403e-05,  3.8604e-03],\n",
+      "        [-1.2524e-02, -1.1022e-02, -1.3438e-02,  ..., -1.4034e-02,\n",
+      "          1.6053e-02,  1.5081e-02]], device='cuda:0'), 'base_model.model.decoder.block.15.layer.1.EncDecAttention.q.lora_A.weight': tensor([[ 0.0230,  0.0099,  0.0328,  ..., -0.0349,  0.0029,  0.0024],\n",
+      "        [ 0.0306,  0.0177, -0.0057,  ..., -0.0538, -0.0250,  0.0284],\n",
+      "        [-0.0387, -0.0495,  0.0077,  ...,  0.0302, -0.0180,  0.0018],\n",
+      "        ...,\n",
+      "        [-0.0281, -0.0137, -0.0024,  ...,  0.0016,  0.0356,  0.0271],\n",
+      "        [-0.0477, -0.0220,  0.0125,  ...,  0.0446,  0.0032, -0.0169],\n",
+      "        [-0.0168, -0.0307, -0.0202,  ...,  0.0207,  0.0262,  0.0194]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.15.layer.1.EncDecAttention.q.lora_B.weight': tensor([[ 0.0084,  0.0045, -0.0095,  ..., -0.0095, -0.0065, -0.0033],\n",
+      "        [ 0.0140,  0.0250, -0.0206,  ..., -0.0109, -0.0197, -0.0172],\n",
+      "        [-0.0027, -0.0009,  0.0136,  ...,  0.0028, -0.0014,  0.0007],\n",
+      "        ...,\n",
+      "        [ 0.0039,  0.0058, -0.0080,  ..., -0.0043, -0.0039, -0.0040],\n",
+      "        [ 0.0005, -0.0047,  0.0135,  ...,  0.0073,  0.0052,  0.0081],\n",
+      "        [-0.0088, -0.0020,  0.0143,  ...,  0.0133,  0.0091,  0.0072]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.15.layer.1.EncDecAttention.v.lora_A.weight': tensor([[-0.0296, -0.0106,  0.0387,  ..., -0.0194, -0.0090, -0.0105],\n",
+      "        [ 0.0402,  0.0177,  0.0116,  ...,  0.0274,  0.0108, -0.0268],\n",
+      "        [-0.0306,  0.0236,  0.0116,  ..., -0.0145,  0.0411,  0.0181],\n",
+      "        ...,\n",
+      "        [-0.0032, -0.0044, -0.0173,  ...,  0.0181, -0.0130, -0.0105],\n",
+      "        [ 0.0389,  0.0156, -0.0165,  ...,  0.0283, -0.0330,  0.0074],\n",
+      "        [-0.0487,  0.0026,  0.0130,  ..., -0.0018,  0.0217, -0.0254]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.15.layer.1.EncDecAttention.v.lora_B.weight': tensor([[-0.0079,  0.0204, -0.0110,  ...,  0.0214,  0.0234, -0.0076],\n",
+      "        [ 0.0076, -0.0051,  0.0005,  ...,  0.0017, -0.0050,  0.0087],\n",
+      "        [-0.0129,  0.0082, -0.0114,  ...,  0.0144,  0.0113, -0.0152],\n",
+      "        ...,\n",
+      "        [-0.0073,  0.0115, -0.0072,  ...,  0.0007,  0.0094, -0.0099],\n",
+      "        [ 0.0005,  0.0024, -0.0036,  ...,  0.0100,  0.0121, -0.0115],\n",
+      "        [-0.0065, -0.0010, -0.0054,  ...,  0.0009,  0.0019, -0.0048]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.16.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0219, -0.0087, -0.0181,  ...,  0.0393,  0.0068, -0.0218],\n",
+      "        [-0.0422, -0.0278,  0.0179,  ...,  0.0154,  0.0026,  0.0087],\n",
+      "        [ 0.0098, -0.0264,  0.0298,  ...,  0.0034, -0.0046, -0.0290],\n",
+      "        ...,\n",
+      "        [ 0.0221, -0.0257,  0.0337,  ...,  0.0161, -0.0059, -0.0238],\n",
+      "        [-0.0034, -0.0086, -0.0082,  ..., -0.0122, -0.0405,  0.0123],\n",
+      "        [ 0.0056, -0.0267, -0.0337,  ...,  0.0398, -0.0013,  0.0338]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.16.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0045,  0.0062,  0.0051,  ..., -0.0066, -0.0064,  0.0069],\n",
+      "        [-0.0070, -0.0047, -0.0063,  ...,  0.0037,  0.0037, -0.0014],\n",
+      "        [ 0.0058,  0.0074, -0.0024,  ..., -0.0094, -0.0086,  0.0102],\n",
+      "        ...,\n",
+      "        [-0.0225, -0.0194,  0.0190,  ...,  0.0209,  0.0221, -0.0232],\n",
+      "        [-0.0137, -0.0164,  0.0105,  ...,  0.0161,  0.0164, -0.0072],\n",
+      "        [ 0.0036,  0.0069,  0.0028,  ..., -0.0018, -0.0002, -0.0341]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.16.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0219, -0.0161,  0.0182,  ...,  0.0006,  0.0182,  0.0124],\n",
+      "        [ 0.0040,  0.0343,  0.0194,  ..., -0.0299, -0.0435,  0.0343],\n",
+      "        [ 0.0039, -0.0184, -0.0296,  ..., -0.0075,  0.0109, -0.0545],\n",
+      "        ...,\n",
+      "        [ 0.0032,  0.0418, -0.0225,  ..., -0.0043, -0.0383, -0.0002],\n",
+      "        [ 0.0360,  0.0251,  0.0007,  ..., -0.0507, -0.0307,  0.0645],\n",
+      "        [ 0.0341,  0.0468,  0.0273,  ..., -0.0019, -0.0207,  0.0422]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.16.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0233, -0.0240,  0.0242,  ..., -0.0231, -0.0250, -0.0250],\n",
+      "        [-0.0088, -0.0091,  0.0081,  ..., -0.0089, -0.0099, -0.0091],\n",
+      "        [ 0.0041,  0.0020,  0.0018,  ...,  0.0014,  0.0007,  0.0016],\n",
+      "        ...,\n",
+      "        [-0.0078, -0.0184,  0.0138,  ..., -0.0104, -0.0209, -0.0149],\n",
+      "        [ 0.0161,  0.0129, -0.0111,  ...,  0.0140,  0.0109,  0.0143],\n",
+      "        [ 0.0067,  0.0030, -0.0056,  ...,  0.0036, -0.0010,  0.0049]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.16.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0027, -0.0303,  0.0067,  ..., -0.0229,  0.0209, -0.0409],\n",
+      "        [-0.0099, -0.0124,  0.0144,  ..., -0.0010, -0.0295,  0.0039],\n",
+      "        [-0.0113,  0.0164,  0.0156,  ...,  0.0317,  0.0263,  0.0163],\n",
+      "        ...,\n",
+      "        [ 0.0317, -0.0133,  0.0245,  ..., -0.0112,  0.0027,  0.0285],\n",
+      "        [ 0.0233,  0.0047, -0.0222,  ..., -0.0179,  0.0321,  0.0024],\n",
+      "        [ 0.0352, -0.0024, -0.0367,  ...,  0.0203, -0.0148, -0.0224]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.16.layer.1.EncDecAttention.q.lora_B.weight': tensor([[ 0.0052, -0.0046,  0.0103,  ...,  0.0058,  0.0120,  0.0090],\n",
+      "        [-0.0022,  0.0010, -0.0026,  ..., -0.0087, -0.0056, -0.0035],\n",
+      "        [ 0.0039, -0.0125,  0.0089,  ..., -0.0214,  0.0074,  0.0066],\n",
+      "        ...,\n",
+      "        [ 0.0275, -0.0290,  0.0311,  ..., -0.0106,  0.0308,  0.0293],\n",
+      "        [ 0.0141, -0.0132,  0.0107,  ..., -0.0075,  0.0121,  0.0140],\n",
+      "        [-0.0304,  0.0323, -0.0312,  ...,  0.0101, -0.0350, -0.0371]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.16.layer.1.EncDecAttention.v.lora_A.weight': tensor([[-0.0211, -0.0161,  0.0199,  ..., -0.0291,  0.0098,  0.0633],\n",
+      "        [ 0.0079, -0.0033, -0.0019,  ...,  0.0192, -0.0015, -0.0375],\n",
+      "        [-0.0076, -0.0288,  0.0126,  ...,  0.0054, -0.0087,  0.0139],\n",
+      "        ...,\n",
+      "        [ 0.0009, -0.0255,  0.0367,  ..., -0.0258,  0.0083,  0.0221],\n",
+      "        [ 0.0277,  0.0011,  0.0027,  ..., -0.0070,  0.0137,  0.0379],\n",
+      "        [ 0.0215,  0.0060, -0.0134,  ...,  0.0053, -0.0184, -0.0192]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.16.layer.1.EncDecAttention.v.lora_B.weight': tensor([[ 0.0108,  0.0078,  0.0119,  ...,  0.0124,  0.0036, -0.0123],\n",
+      "        [ 0.0202, -0.0023,  0.0220,  ...,  0.0209, -0.0016, -0.0214],\n",
+      "        [-0.0082,  0.0027,  0.0075,  ..., -0.0011, -0.0155,  0.0050],\n",
+      "        ...,\n",
+      "        [-0.0056,  0.0035, -0.0025,  ..., -0.0039, -0.0004,  0.0047],\n",
+      "        [-0.0191,  0.0135, -0.0108,  ..., -0.0092, -0.0125,  0.0128],\n",
+      "        [ 0.0191, -0.0031,  0.0218,  ...,  0.0199,  0.0088, -0.0202]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.17.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0110,  0.0033, -0.0113,  ...,  0.0048,  0.0247,  0.0054],\n",
+      "        [ 0.0127, -0.0212, -0.0472,  ..., -0.0268,  0.0428, -0.0202],\n",
+      "        [-0.0078,  0.0169,  0.0111,  ..., -0.0042, -0.0003, -0.0235],\n",
+      "        ...,\n",
+      "        [-0.0244, -0.0122, -0.0304,  ..., -0.0329,  0.0367, -0.0051],\n",
+      "        [ 0.0205,  0.0283,  0.0346,  ...,  0.0034,  0.0209,  0.0203],\n",
+      "        [-0.0127,  0.0077,  0.0263,  ...,  0.0332, -0.0028,  0.0078]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.17.layer.0.SelfAttention.q.lora_B.weight': tensor([[-0.0027,  0.0025,  0.0065,  ...,  0.0052, -0.0079, -0.0049],\n",
+      "        [ 0.0232,  0.0130,  0.0136,  ...,  0.0119, -0.0129, -0.0109],\n",
+      "        [-0.0266, -0.0229, -0.0219,  ..., -0.0216,  0.0204,  0.0209],\n",
+      "        ...,\n",
+      "        [ 0.0131,  0.0106,  0.0096,  ...,  0.0094, -0.0127, -0.0143],\n",
+      "        [ 0.0107,  0.0033,  0.0019,  ...,  0.0027, -0.0010, -0.0041],\n",
+      "        [-0.0081, -0.0046, -0.0037,  ..., -0.0038,  0.0019,  0.0018]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.17.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0325,  0.0404, -0.0134,  ...,  0.0062, -0.0371,  0.0013],\n",
+      "        [-0.0041, -0.0016,  0.0144,  ..., -0.0214, -0.0108, -0.0138],\n",
+      "        [-0.0005,  0.0322, -0.0368,  ..., -0.0403,  0.0035, -0.0059],\n",
+      "        ...,\n",
+      "        [ 0.0045,  0.0389,  0.0109,  ..., -0.0075, -0.0146,  0.0066],\n",
+      "        [-0.0320, -0.0356,  0.0048,  ...,  0.0524,  0.0445, -0.0348],\n",
+      "        [-0.0180, -0.0509,  0.0261,  ...,  0.0295,  0.0356, -0.0276]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.17.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0078,  0.0004,  0.0029,  ..., -0.0020,  0.0022,  0.0038],\n",
+      "        [ 0.0122,  0.0044,  0.0141,  ...,  0.0164, -0.0155, -0.0095],\n",
+      "        [-0.0067,  0.0061,  0.0006,  ..., -0.0035,  0.0040,  0.0004],\n",
+      "        ...,\n",
+      "        [ 0.0147,  0.0099,  0.0087,  ...,  0.0133, -0.0135, -0.0156],\n",
+      "        [ 0.0059, -0.0010,  0.0012,  ...,  0.0056, -0.0035,  0.0037],\n",
+      "        [ 0.0002,  0.0016,  0.0042,  ...,  0.0002,  0.0006,  0.0034]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.17.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0348, -0.0118,  0.0138,  ..., -0.0390,  0.0242, -0.0179],\n",
+      "        [ 0.0164, -0.0063,  0.0224,  ..., -0.0279,  0.0316, -0.0096],\n",
+      "        [ 0.0018, -0.0266,  0.0202,  ..., -0.0475, -0.0246, -0.0220],\n",
+      "        ...,\n",
+      "        [ 0.0014, -0.0097, -0.0064,  ...,  0.0311,  0.0035, -0.0014],\n",
+      "        [ 0.0395,  0.0291,  0.0073,  ..., -0.0023, -0.0247,  0.0419],\n",
+      "        [ 0.0226,  0.0164, -0.0280,  ...,  0.0056,  0.0016,  0.0153]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.17.layer.1.EncDecAttention.q.lora_B.weight': tensor([[ 0.0078,  0.0093,  0.0078,  ..., -0.0093, -0.0051, -0.0021],\n",
+      "        [-0.0284, -0.0147, -0.0206,  ...,  0.0180,  0.0384,  0.0393],\n",
+      "        [ 0.0054,  0.0074,  0.0035,  ..., -0.0096, -0.0074, -0.0064],\n",
+      "        ...,\n",
+      "        [-0.0116, -0.0127, -0.0108,  ...,  0.0129,  0.0128,  0.0102],\n",
+      "        [ 0.0139,  0.0120,  0.0094,  ..., -0.0132, -0.0130, -0.0141],\n",
+      "        [ 0.0109,  0.0123,  0.0065,  ..., -0.0120, -0.0124, -0.0094]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.17.layer.1.EncDecAttention.v.lora_A.weight': tensor([[-0.0262, -0.0628,  0.0287,  ..., -0.0263,  0.0311, -0.0192],\n",
+      "        [-0.0029,  0.0124, -0.0072,  ...,  0.0043, -0.0331, -0.0429],\n",
+      "        [ 0.0236,  0.0187, -0.0174,  ...,  0.0047,  0.0140,  0.0135],\n",
+      "        ...,\n",
+      "        [ 0.0152,  0.0109,  0.0131,  ..., -0.0012, -0.0165,  0.0063],\n",
+      "        [-0.0361, -0.0315,  0.0343,  ..., -0.0198,  0.0440, -0.0010],\n",
+      "        [ 0.0114,  0.0118, -0.0288,  ...,  0.0051, -0.0095, -0.0169]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.17.layer.1.EncDecAttention.v.lora_B.weight': tensor([[-0.0318,  0.0044,  0.0069,  ...,  0.0107, -0.0189,  0.0074],\n",
+      "        [-0.0332,  0.0219, -0.0159,  ..., -0.0156, -0.0224,  0.0240],\n",
+      "        [ 0.0288, -0.0129,  0.0086,  ...,  0.0045,  0.0240, -0.0148],\n",
+      "        ...,\n",
+      "        [ 0.0157,  0.0096, -0.0117,  ..., -0.0127, -0.0148, -0.0194],\n",
+      "        [ 0.0155,  0.0098, -0.0099,  ..., -0.0118, -0.0114, -0.0147],\n",
+      "        [-0.0018, -0.0110,  0.0131,  ...,  0.0123,  0.0138,  0.0039]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.18.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0352, -0.0196,  0.0188,  ...,  0.0375, -0.0202, -0.0196],\n",
+      "        [ 0.0414,  0.0130, -0.0185,  ...,  0.0279, -0.0007, -0.0084],\n",
+      "        [ 0.0244, -0.0149, -0.0413,  ...,  0.0338, -0.0023,  0.0292],\n",
+      "        ...,\n",
+      "        [ 0.0105,  0.0305,  0.0255,  ..., -0.0105, -0.0306, -0.0135],\n",
+      "        [ 0.0201,  0.0205,  0.0201,  ...,  0.0055, -0.0065, -0.0440],\n",
+      "        [ 0.0362, -0.0236,  0.0030,  ..., -0.0126,  0.0508,  0.0348]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.18.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0081,  0.0077,  0.0071,  ..., -0.0068,  0.0024,  0.0021],\n",
+      "        [-0.0012,  0.0028,  0.0023,  ..., -0.0028, -0.0080,  0.0022],\n",
+      "        [ 0.0074,  0.0088,  0.0088,  ..., -0.0086,  0.0027,  0.0098],\n",
+      "        ...,\n",
+      "        [-0.0137, -0.0088, -0.0094,  ...,  0.0106,  0.0117, -0.0056],\n",
+      "        [ 0.0091,  0.0049,  0.0071,  ..., -0.0084,  0.0008, -0.0027],\n",
+      "        [-0.0011, -0.0091, -0.0071,  ...,  0.0047, -0.0046, -0.0082]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.18.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0202,  0.0472, -0.0144,  ..., -0.0219,  0.0103,  0.0043],\n",
+      "        [ 0.0089,  0.0365,  0.0118,  ..., -0.0073,  0.0052,  0.0454],\n",
+      "        [-0.0188,  0.0213, -0.0220,  ..., -0.0063, -0.0301,  0.0115],\n",
+      "        ...,\n",
+      "        [ 0.0001,  0.0202,  0.0150,  ..., -0.0270,  0.0132,  0.0644],\n",
+      "        [-0.0128,  0.0170,  0.0343,  ...,  0.0307,  0.0125,  0.0154],\n",
+      "        [ 0.0368, -0.0510, -0.0221,  ...,  0.0050,  0.0360, -0.0401]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.18.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0183,  0.0209,  0.0223,  ...,  0.0224, -0.0276, -0.0238],\n",
+      "        [ 0.0018,  0.0003, -0.0033,  ..., -0.0021,  0.0007,  0.0008],\n",
+      "        [-0.0094, -0.0069, -0.0029,  ..., -0.0034,  0.0021,  0.0041],\n",
+      "        ...,\n",
+      "        [ 0.0106,  0.0082, -0.0028,  ..., -0.0003,  0.0051,  0.0002],\n",
+      "        [ 0.0205,  0.0214,  0.0110,  ...,  0.0128, -0.0095, -0.0170],\n",
+      "        [-0.0337, -0.0304, -0.0236,  ..., -0.0311,  0.0311,  0.0223]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.18.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0321, -0.0146,  0.0019,  ...,  0.0236,  0.0183, -0.0471],\n",
+      "        [-0.0162, -0.0074, -0.0296,  ...,  0.0163, -0.0134, -0.0315],\n",
+      "        [ 0.0442, -0.0036, -0.0010,  ..., -0.0262,  0.0005,  0.0129],\n",
+      "        ...,\n",
+      "        [ 0.0039,  0.0022, -0.0167,  ..., -0.0317,  0.0070,  0.0184],\n",
+      "        [-0.0482, -0.0197, -0.0316,  ..., -0.0014, -0.0049, -0.0271],\n",
+      "        [ 0.0432, -0.0289,  0.0295,  ..., -0.0390, -0.0308,  0.0124]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.18.layer.1.EncDecAttention.q.lora_B.weight': tensor([[-0.0075, -0.0108,  0.0096,  ...,  0.0016, -0.0095,  0.0080],\n",
+      "        [-0.0140, -0.0120,  0.0073,  ...,  0.0127, -0.0119,  0.0123],\n",
+      "        [-0.0056, -0.0077,  0.0043,  ..., -0.0065, -0.0042,  0.0043],\n",
+      "        ...,\n",
+      "        [-0.0019,  0.0033,  0.0096,  ..., -0.0167, -0.0077,  0.0028],\n",
+      "        [-0.0132, -0.0107,  0.0170,  ...,  0.0081, -0.0119,  0.0058],\n",
+      "        [-0.0042, -0.0003,  0.0090,  ...,  0.0026, -0.0019,  0.0004]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.18.layer.1.EncDecAttention.v.lora_A.weight': tensor([[ 0.0196,  0.0030,  0.0224,  ..., -0.0038,  0.0045,  0.0215],\n",
+      "        [ 0.0100,  0.0074,  0.0359,  ...,  0.0210,  0.0242,  0.0113],\n",
+      "        [ 0.0026, -0.0086,  0.0127,  ...,  0.0132, -0.0101,  0.0025],\n",
+      "        ...,\n",
+      "        [ 0.0285, -0.0291,  0.0067,  ..., -0.0070, -0.0483, -0.0456],\n",
+      "        [-0.0019, -0.0021, -0.0288,  ..., -0.0177,  0.0375,  0.0036],\n",
+      "        [-0.0219,  0.0325,  0.0195,  ..., -0.0087, -0.0078,  0.0244]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.18.layer.1.EncDecAttention.v.lora_B.weight': tensor([[ 0.0184, -0.0062, -0.0195,  ..., -0.0182,  0.0185,  0.0185],\n",
+      "        [ 0.0150, -0.0073, -0.0114,  ..., -0.0022,  0.0037,  0.0024],\n",
+      "        [ 0.0261, -0.0178, -0.0264,  ..., -0.0250,  0.0274,  0.0250],\n",
+      "        ...,\n",
+      "        [ 0.0242, -0.0040, -0.0231,  ..., -0.0215,  0.0186,  0.0224],\n",
+      "        [-0.0210,  0.0187,  0.0201,  ...,  0.0210, -0.0218, -0.0205],\n",
+      "        [-0.0242, -0.0048,  0.0228,  ...,  0.0187, -0.0221, -0.0207]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.19.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0391, -0.0192,  0.0063,  ...,  0.0308,  0.0023,  0.0011],\n",
+      "        [-0.0155,  0.0229, -0.0224,  ..., -0.0024,  0.0059, -0.0193],\n",
+      "        [ 0.0280, -0.0313, -0.0107,  ..., -0.0039,  0.0153,  0.0012],\n",
+      "        ...,\n",
+      "        [ 0.0262,  0.0110, -0.0075,  ..., -0.0037,  0.0194, -0.0427],\n",
+      "        [-0.0242, -0.0382,  0.0120,  ...,  0.0300, -0.0108,  0.0073],\n",
+      "        [-0.0040, -0.0096,  0.0203,  ..., -0.0097, -0.0198, -0.0229]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.19.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 6.1453e-03,  6.1973e-03, -6.8584e-03,  ..., -4.5252e-03,\n",
+      "          6.3516e-03,  6.2948e-03],\n",
+      "        [-4.5412e-03, -4.4871e-03,  4.0846e-03,  ...,  5.9160e-03,\n",
+      "         -4.4922e-03, -4.3107e-03],\n",
+      "        [-6.2454e-03, -6.1674e-03,  6.7055e-03,  ...,  4.5486e-03,\n",
+      "         -6.3253e-03, -6.0943e-03],\n",
+      "        ...,\n",
+      "        [-6.4837e-03, -7.3207e-03,  8.4997e-03,  ..., -6.3963e-05,\n",
+      "         -3.4027e-04,  1.2465e-02],\n",
+      "        [ 1.0956e-02,  9.6466e-03, -7.7077e-03,  ..., -9.7397e-03,\n",
+      "          1.1114e-02, -2.4019e-02],\n",
+      "        [ 2.1932e-02,  1.3618e-02, -1.6683e-02,  ..., -1.7114e-02,\n",
+      "          1.7331e-02, -2.1152e-02]], device='cuda:0'), 'base_model.model.decoder.block.19.layer.0.SelfAttention.v.lora_A.weight': tensor([[-1.2191e-03, -1.8192e-02,  2.0656e-02,  ..., -3.0691e-02,\n",
+      "         -3.7022e-02,  6.2503e-02],\n",
+      "        [-8.9603e-05, -2.4246e-02, -1.6662e-02,  ..., -4.4149e-02,\n",
+      "         -7.4380e-03,  1.4131e-02],\n",
+      "        [-1.2862e-03,  2.0751e-02,  1.4463e-03,  ...,  1.3739e-02,\n",
+      "          3.9745e-02, -2.3024e-02],\n",
+      "        ...,\n",
+      "        [ 1.6170e-03,  2.8526e-02, -1.6437e-02,  ..., -4.2300e-02,\n",
+      "         -1.4941e-02,  4.1151e-02],\n",
+      "        [-1.9709e-02, -1.1424e-02,  1.7977e-02,  ...,  3.2824e-02,\n",
+      "          6.1655e-03, -5.5113e-02],\n",
+      "        [-4.6231e-03, -1.9727e-03,  2.2364e-02,  ..., -3.3645e-03,\n",
+      "         -1.6693e-02, -9.5787e-04]], device='cuda:0'), 'base_model.model.decoder.block.19.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0227,  0.0295, -0.0247,  ...,  0.0254, -0.0258,  0.0219],\n",
+      "        [ 0.0271,  0.0284, -0.0212,  ...,  0.0271, -0.0269,  0.0150],\n",
+      "        [ 0.0234,  0.0235, -0.0208,  ...,  0.0246, -0.0221,  0.0152],\n",
+      "        ...,\n",
+      "        [ 0.0078,  0.0068, -0.0066,  ...,  0.0072, -0.0073,  0.0119],\n",
+      "        [ 0.0140,  0.0122, -0.0158,  ...,  0.0121, -0.0126,  0.0089],\n",
+      "        [-0.0119, -0.0147,  0.0114,  ..., -0.0144,  0.0136, -0.0122]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.19.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0183,  0.0285, -0.0087,  ...,  0.0389, -0.0304, -0.0211],\n",
+      "        [-0.0469,  0.0042, -0.0137,  ...,  0.0184, -0.0258,  0.0023],\n",
+      "        [ 0.0009, -0.0048, -0.0200,  ...,  0.0073, -0.0202,  0.0201],\n",
+      "        ...,\n",
+      "        [-0.0237,  0.0142,  0.0049,  ...,  0.0029, -0.0308,  0.0134],\n",
+      "        [-0.0482,  0.0319, -0.0207,  ...,  0.0067, -0.0235, -0.0195],\n",
+      "        [ 0.0306, -0.0269,  0.0360,  ..., -0.0022, -0.0164,  0.0399]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.19.layer.1.EncDecAttention.q.lora_B.weight': tensor([[-0.0137, -0.0193,  0.0165,  ..., -0.0196, -0.0190,  0.0197],\n",
+      "        [ 0.0025, -0.0017,  0.0025,  ..., -0.0013, -0.0026,  0.0021],\n",
+      "        [-0.0140, -0.0209,  0.0177,  ..., -0.0156, -0.0172,  0.0192],\n",
+      "        ...,\n",
+      "        [ 0.0085,  0.0040, -0.0062,  ...,  0.0022,  0.0034, -0.0028],\n",
+      "        [-0.0045, -0.0048,  0.0089,  ..., -0.0104, -0.0076,  0.0103],\n",
+      "        [-0.0155, -0.0158,  0.0103,  ..., -0.0147, -0.0128,  0.0156]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.19.layer.1.EncDecAttention.v.lora_A.weight': tensor([[-0.0591, -0.0330,  0.0479,  ..., -0.0135,  0.0269, -0.0320],\n",
+      "        [ 0.0116,  0.0332,  0.0454,  ..., -0.0533,  0.0502,  0.0280],\n",
+      "        [ 0.0680,  0.0278, -0.0021,  ...,  0.0058, -0.0372,  0.0176],\n",
+      "        ...,\n",
+      "        [-0.0592, -0.0317,  0.0202,  ...,  0.0044,  0.0342, -0.0396],\n",
+      "        [-0.0679, -0.0569,  0.0012,  ...,  0.0163,  0.0423, -0.0169],\n",
+      "        [ 0.0339,  0.0452, -0.0357,  ...,  0.0423, -0.0517,  0.0218]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.19.layer.1.EncDecAttention.v.lora_B.weight': tensor([[-0.0204, -0.0194, -0.0167,  ..., -0.0192,  0.0034,  0.0193],\n",
+      "        [-0.0046, -0.0046, -0.0007,  ..., -0.0025,  0.0107,  0.0048],\n",
+      "        [-0.0109, -0.0102, -0.0077,  ..., -0.0106,  0.0037,  0.0111],\n",
+      "        ...,\n",
+      "        [ 0.0105, -0.0051, -0.0013,  ..., -0.0022, -0.0002,  0.0012],\n",
+      "        [ 0.0298,  0.0384, -0.0381,  ...,  0.0380,  0.0370, -0.0360],\n",
+      "        [-0.0161, -0.0247,  0.0242,  ..., -0.0196, -0.0242,  0.0225]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.20.layer.0.SelfAttention.q.lora_A.weight': tensor([[-0.0330, -0.0107,  0.0211,  ...,  0.0107,  0.0013, -0.0105],\n",
+      "        [ 0.0142, -0.0134,  0.0259,  ...,  0.0185,  0.0278,  0.0195],\n",
+      "        [-0.0087, -0.0145, -0.0396,  ...,  0.0157, -0.0085, -0.0097],\n",
+      "        ...,\n",
+      "        [-0.0029,  0.0066, -0.0174,  ...,  0.0006, -0.0260, -0.0203],\n",
+      "        [ 0.0232,  0.0241,  0.0129,  ..., -0.0420,  0.0071, -0.0258],\n",
+      "        [ 0.0318,  0.0162,  0.0062,  ...,  0.0249,  0.0159, -0.0238]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.20.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 0.0066,  0.0075,  0.0076,  ..., -0.0074, -0.0073, -0.0012],\n",
+      "        [-0.0060, -0.0079, -0.0063,  ...,  0.0062,  0.0067,  0.0021],\n",
+      "        [-0.0026, -0.0072, -0.0017,  ...,  0.0012,  0.0013, -0.0047],\n",
+      "        ...,\n",
+      "        [-0.0108, -0.0189, -0.0058,  ...,  0.0137,  0.0104,  0.0008],\n",
+      "        [-0.0061, -0.0086, -0.0066,  ...,  0.0030,  0.0012, -0.0078],\n",
+      "        [ 0.0148,  0.0042,  0.0149,  ..., -0.0157, -0.0143, -0.0109]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.20.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0134,  0.0197,  0.0471,  ..., -0.0059, -0.0069,  0.0288],\n",
+      "        [-0.0282, -0.0023, -0.0330,  ...,  0.0061,  0.0268, -0.0604],\n",
+      "        [-0.0080,  0.0056,  0.0321,  ..., -0.0319,  0.0065,  0.0143],\n",
+      "        ...,\n",
+      "        [-0.0269,  0.0138,  0.0644,  ...,  0.0223, -0.0018,  0.0240],\n",
+      "        [-0.0010, -0.0047,  0.0269,  ...,  0.0136,  0.0345, -0.0001],\n",
+      "        [-0.0323,  0.0287,  0.0055,  ...,  0.0195,  0.0189, -0.0179]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.20.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0061,  0.0025, -0.0070,  ...,  0.0109,  0.0056,  0.0067],\n",
+      "        [-0.0246,  0.0245, -0.0271,  ..., -0.0096,  0.0270,  0.0278],\n",
+      "        [ 0.0282, -0.0249,  0.0236,  ...,  0.0410, -0.0145, -0.0263],\n",
+      "        ...,\n",
+      "        [-0.0128,  0.0086,  0.0037,  ..., -0.0332, -0.0036,  0.0066],\n",
+      "        [-0.0181,  0.0180, -0.0174,  ..., -0.0017,  0.0149,  0.0184],\n",
+      "        [-0.0168,  0.0144, -0.0048,  ...,  0.0045,  0.0113,  0.0176]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.20.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0309, -0.0433, -0.0072,  ...,  0.0526, -0.0267,  0.0111],\n",
+      "        [ 0.0322,  0.0130, -0.0075,  ..., -0.0538,  0.0289,  0.0092],\n",
+      "        [ 0.0119,  0.0142,  0.0159,  ..., -0.0186,  0.0124,  0.0052],\n",
+      "        ...,\n",
+      "        [-0.0182, -0.0393,  0.0251,  ...,  0.0051, -0.0294,  0.0166],\n",
+      "        [ 0.0204, -0.0257,  0.0171,  ...,  0.0160,  0.0026,  0.0258],\n",
+      "        [-0.0115, -0.0061,  0.0026,  ..., -0.0053,  0.0131,  0.0137]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.20.layer.1.EncDecAttention.q.lora_B.weight': tensor([[ 0.0007, -0.0135, -0.0022,  ..., -0.0014,  0.0102, -0.0091],\n",
+      "        [-0.0013, -0.0089, -0.0008,  ..., -0.0011,  0.0063, -0.0077],\n",
+      "        [ 0.0135, -0.0018, -0.0098,  ...,  0.0087,  0.0182, -0.0043],\n",
+      "        ...,\n",
+      "        [ 0.0196, -0.0158, -0.0186,  ...,  0.0112,  0.0094, -0.0163],\n",
+      "        [ 0.0239, -0.0141, -0.0202,  ...,  0.0215,  0.0194, -0.0229],\n",
+      "        [-0.0105,  0.0090,  0.0092,  ..., -0.0076, -0.0030,  0.0103]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.20.layer.1.EncDecAttention.v.lora_A.weight': tensor([[ 0.0300, -0.0317, -0.0016,  ...,  0.0477, -0.0341, -0.0443],\n",
+      "        [ 0.0130, -0.0124,  0.0146,  ...,  0.0177,  0.0036,  0.0013],\n",
+      "        [-0.0037, -0.0089, -0.0270,  ..., -0.0676,  0.0243,  0.0449],\n",
+      "        ...,\n",
+      "        [ 0.0045,  0.0037, -0.0003,  ..., -0.0588,  0.0432,  0.0416],\n",
+      "        [ 0.0265, -0.0409, -0.0213,  ...,  0.0781, -0.0029, -0.0608],\n",
+      "        [-0.0193, -0.0034,  0.0216,  ..., -0.0217,  0.0456,  0.0479]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.20.layer.1.EncDecAttention.v.lora_B.weight': tensor([[-0.0056, -0.0064,  0.0067,  ..., -0.0014, -0.0009,  0.0006],\n",
+      "        [-0.0254, -0.0245,  0.0249,  ...,  0.0180, -0.0263,  0.0243],\n",
+      "        [-0.0236, -0.0116,  0.0207,  ...,  0.0163, -0.0246,  0.0150],\n",
+      "        ...,\n",
+      "        [ 0.0120,  0.0035, -0.0083,  ..., -0.0079,  0.0130, -0.0055],\n",
+      "        [ 0.0174,  0.0187, -0.0193,  ..., -0.0140,  0.0158, -0.0161],\n",
+      "        [ 0.0092,  0.0099, -0.0160,  ..., -0.0007,  0.0068, -0.0113]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.21.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0343,  0.0204,  0.0275,  ...,  0.0053, -0.0517, -0.0413],\n",
+      "        [-0.0248, -0.0211,  0.0076,  ...,  0.0244,  0.0250,  0.0150],\n",
+      "        [-0.0068, -0.0234,  0.0406,  ...,  0.0005, -0.0413, -0.0170],\n",
+      "        ...,\n",
+      "        [ 0.0428, -0.0168, -0.0036,  ..., -0.0061, -0.0196,  0.0055],\n",
+      "        [-0.0289,  0.0417, -0.0210,  ...,  0.0379,  0.0039,  0.0105],\n",
+      "        [ 0.0070,  0.0048,  0.0043,  ..., -0.0240, -0.0300, -0.0242]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.21.layer.0.SelfAttention.q.lora_B.weight': tensor([[ 4.4348e-03, -2.6485e-03,  2.1192e-03,  ...,  6.5235e-04,\n",
+      "         -4.2449e-03,  2.7613e-03],\n",
+      "        [ 1.5684e-03, -2.8341e-04,  4.5689e-05,  ..., -1.3167e-03,\n",
+      "         -6.2685e-04, -4.6340e-04],\n",
+      "        [ 5.4143e-04,  6.6562e-04, -9.4752e-04,  ..., -2.5828e-03,\n",
+      "          7.1928e-04, -1.4943e-03],\n",
+      "        ...,\n",
+      "        [ 3.1052e-03, -1.0300e-02,  9.1300e-04,  ...,  5.9527e-03,\n",
+      "          1.4728e-03,  2.2390e-03],\n",
+      "        [-1.0343e-02,  3.8973e-03, -1.0227e-02,  ..., -2.1103e-03,\n",
+      "          1.1451e-02, -1.0909e-02],\n",
+      "        [ 9.7546e-03, -4.5729e-03,  9.8230e-03,  ...,  6.4971e-03,\n",
+      "         -9.7390e-03,  9.7719e-03]], device='cuda:0'), 'base_model.model.decoder.block.21.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0403, -0.0040, -0.0013,  ..., -0.0024,  0.0256,  0.0428],\n",
+      "        [-0.0482,  0.0095, -0.0140,  ..., -0.0013,  0.0166, -0.0315],\n",
+      "        [ 0.0045,  0.0172,  0.0354,  ..., -0.0425, -0.0294,  0.0075],\n",
+      "        ...,\n",
+      "        [-0.0317, -0.0092, -0.0335,  ...,  0.0219,  0.0336, -0.0158],\n",
+      "        [-0.0078, -0.0031, -0.0020,  ..., -0.0177, -0.0036, -0.0421],\n",
+      "        [ 0.0143, -0.0200, -0.0190,  ...,  0.0009, -0.0299,  0.0199]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.21.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0263, -0.0255,  0.0235,  ..., -0.0238, -0.0280,  0.0286],\n",
+      "        [ 0.0201, -0.0229,  0.0209,  ..., -0.0200, -0.0184,  0.0127],\n",
+      "        [-0.0039,  0.0015, -0.0030,  ..., -0.0007,  0.0009, -0.0090],\n",
+      "        ...,\n",
+      "        [ 0.0110, -0.0022,  0.0084,  ..., -0.0041,  0.0051,  0.0059],\n",
+      "        [ 0.0110, -0.0111,  0.0048,  ..., -0.0118, -0.0143,  0.0013],\n",
+      "        [ 0.0178, -0.0232,  0.0145,  ..., -0.0203, -0.0183,  0.0250]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.21.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0286, -0.0409, -0.0219,  ...,  0.0076,  0.0002, -0.0041],\n",
+      "        [-0.0026,  0.0309,  0.0021,  ..., -0.0009, -0.0211,  0.0268],\n",
+      "        [ 0.0109,  0.0248, -0.0301,  ...,  0.0485, -0.0336, -0.0063],\n",
+      "        ...,\n",
+      "        [-0.0150, -0.0146, -0.0259,  ..., -0.0196,  0.0254, -0.0009],\n",
+      "        [ 0.0092, -0.0122, -0.0444,  ...,  0.0018, -0.0069,  0.0313],\n",
+      "        [-0.0075, -0.0570, -0.0182,  ...,  0.0037,  0.0248, -0.0176]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.21.layer.1.EncDecAttention.q.lora_B.weight': tensor([[-0.0168,  0.0255,  0.0151,  ..., -0.0183, -0.0177, -0.0194],\n",
+      "        [-0.0206,  0.0197,  0.0183,  ..., -0.0195, -0.0225, -0.0275],\n",
+      "        [ 0.0361, -0.0352, -0.0348,  ...,  0.0332,  0.0327,  0.0342],\n",
+      "        ...,\n",
+      "        [ 0.0194, -0.0234, -0.0214,  ...,  0.0162,  0.0178,  0.0205],\n",
+      "        [-0.0017,  0.0059,  0.0040,  ..., -0.0048, -0.0056, -0.0058],\n",
+      "        [ 0.0077, -0.0114, -0.0011,  ...,  0.0101,  0.0192,  0.0077]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.21.layer.1.EncDecAttention.v.lora_A.weight': tensor([[-3.5695e-02,  1.6790e-02,  5.7587e-04,  ..., -1.6495e-02,\n",
+      "          1.2606e-02,  4.8337e-02],\n",
+      "        [ 2.3041e-02, -3.3336e-02, -2.2698e-02,  ...,  5.9429e-02,\n",
+      "         -3.1289e-02, -4.4992e-02],\n",
+      "        [-1.5568e-02, -5.2893e-02, -1.2566e-02,  ...,  4.0882e-02,\n",
+      "          1.0086e-02, -9.9746e-03],\n",
+      "        ...,\n",
+      "        [ 2.4377e-02, -3.7222e-02,  1.3223e-02,  ...,  5.4595e-02,\n",
+      "         -3.3970e-02, -1.3829e-03],\n",
+      "        [ 1.6691e-02,  5.6453e-02, -1.8042e-02,  ..., -7.6591e-02,\n",
+      "          9.8997e-05,  4.6439e-02],\n",
+      "        [-9.6895e-03, -2.9741e-02,  1.2167e-02,  ...,  2.4567e-02,\n",
+      "          9.3858e-03, -3.2569e-02]], device='cuda:0'), 'base_model.model.decoder.block.21.layer.1.EncDecAttention.v.lora_B.weight': tensor([[ 0.0311, -0.0313, -0.0311,  ..., -0.0332,  0.0247, -0.0304],\n",
+      "        [ 0.0029,  0.0026,  0.0029,  ...,  0.0035, -0.0013, -0.0060],\n",
+      "        [-0.0258,  0.0229,  0.0264,  ...,  0.0237, -0.0149,  0.0358],\n",
+      "        ...,\n",
+      "        [ 0.0278, -0.0260, -0.0260,  ..., -0.0240,  0.0291, -0.0316],\n",
+      "        [ 0.0034,  0.0013, -0.0023,  ..., -0.0007, -0.0021, -0.0052],\n",
+      "        [ 0.0103, -0.0093, -0.0117,  ..., -0.0091,  0.0129, -0.0151]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.22.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0075,  0.0101,  0.0066,  ...,  0.0137, -0.0290, -0.0332],\n",
+      "        [-0.0080, -0.0053, -0.0135,  ...,  0.0092, -0.0176, -0.0109],\n",
+      "        [-0.0237, -0.0107,  0.0388,  ..., -0.0160, -0.0025, -0.0372],\n",
+      "        ...,\n",
+      "        [ 0.0404, -0.0182,  0.0073,  ...,  0.0007,  0.0072, -0.0113],\n",
+      "        [ 0.0124, -0.0111,  0.0153,  ..., -0.0410, -0.0116, -0.0214],\n",
+      "        [-0.0053, -0.0287, -0.0165,  ..., -0.0356,  0.0135, -0.0157]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.22.layer.0.SelfAttention.q.lora_B.weight': tensor([[-0.0063,  0.0048,  0.0027,  ...,  0.0076, -0.0064, -0.0074],\n",
+      "        [ 0.0014, -0.0117, -0.0015,  ..., -0.0004,  0.0117,  0.0018],\n",
+      "        [ 0.0156, -0.0110,  0.0123,  ..., -0.0135,  0.0156,  0.0141],\n",
+      "        ...,\n",
+      "        [ 0.0047, -0.0093,  0.0169,  ..., -0.0119,  0.0120,  0.0156],\n",
+      "        [-0.0021,  0.0090, -0.0040,  ...,  0.0050, -0.0092, -0.0039],\n",
+      "        [-0.0010, -0.0057, -0.0066,  ..., -0.0030,  0.0066, -0.0035]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.22.layer.0.SelfAttention.v.lora_A.weight': tensor([[ 0.0456, -0.0016, -0.0174,  ...,  0.0192, -0.0071, -0.0623],\n",
+      "        [ 0.0340,  0.0215, -0.0414,  ..., -0.0354, -0.0016, -0.0324],\n",
+      "        [ 0.0472, -0.0280, -0.0585,  ..., -0.0303, -0.0207, -0.0495],\n",
+      "        ...,\n",
+      "        [-0.0485, -0.0319, -0.0091,  ...,  0.0348,  0.0049,  0.0335],\n",
+      "        [-0.0007,  0.0192, -0.0118,  ..., -0.0202, -0.0005, -0.0519],\n",
+      "        [-0.0050, -0.0187,  0.0160,  ...,  0.0365,  0.0291,  0.0377]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.22.layer.0.SelfAttention.v.lora_B.weight': tensor([[-0.0233, -0.0288, -0.0115,  ...,  0.0231, -0.0130,  0.0227],\n",
+      "        [-0.0201, -0.0285, -0.0253,  ...,  0.0275, -0.0260,  0.0299],\n",
+      "        [-0.0058, -0.0212, -0.0120,  ...,  0.0103, -0.0118,  0.0128],\n",
+      "        ...,\n",
+      "        [ 0.0127,  0.0124,  0.0190,  ..., -0.0125,  0.0157, -0.0204],\n",
+      "        [-0.0180, -0.0212, -0.0175,  ...,  0.0185, -0.0135,  0.0195],\n",
+      "        [-0.0200, -0.0104, -0.0042,  ...,  0.0166, -0.0149,  0.0107]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.22.layer.1.EncDecAttention.q.lora_A.weight': tensor([[-0.0084,  0.0150, -0.0145,  ..., -0.0232, -0.0257, -0.0149],\n",
+      "        [ 0.0048, -0.0269,  0.0020,  ..., -0.0380, -0.0280,  0.0107],\n",
+      "        [-0.0511,  0.0117, -0.0309,  ..., -0.0178, -0.0143,  0.0124],\n",
+      "        ...,\n",
+      "        [ 0.0064,  0.0060,  0.0122,  ...,  0.0043, -0.0113,  0.0389],\n",
+      "        [-0.0138,  0.0269,  0.0226,  ...,  0.0071, -0.0021, -0.0120],\n",
+      "        [-0.0171,  0.0366, -0.0063,  ...,  0.0260,  0.0141, -0.0228]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.22.layer.1.EncDecAttention.q.lora_B.weight': tensor([[ 0.0172,  0.0150, -0.0139,  ..., -0.0159, -0.0151, -0.0160],\n",
+      "        [ 0.0052,  0.0058,  0.0026,  ..., -0.0062, -0.0036, -0.0068],\n",
+      "        [ 0.0252,  0.0213, -0.0235,  ..., -0.0217, -0.0223, -0.0177],\n",
+      "        ...,\n",
+      "        [-0.0061, -0.0087,  0.0110,  ...,  0.0117,  0.0126,  0.0129],\n",
+      "        [-0.0169, -0.0116,  0.0146,  ...,  0.0165,  0.0159,  0.0200],\n",
+      "        [ 0.0153,  0.0157, -0.0118,  ..., -0.0148, -0.0138, -0.0143]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.22.layer.1.EncDecAttention.v.lora_A.weight': tensor([[ 0.0323,  0.0615, -0.0286,  ..., -0.0212,  0.0202,  0.0314],\n",
+      "        [-0.0259,  0.0106, -0.0415,  ..., -0.0654, -0.0317,  0.0154],\n",
+      "        [-0.0052, -0.0413,  0.0448,  ..., -0.0346,  0.0133, -0.0400],\n",
+      "        ...,\n",
+      "        [ 0.0140,  0.0323, -0.0152,  ..., -0.0569, -0.0013,  0.0601],\n",
+      "        [-0.0142, -0.0644,  0.0147,  ...,  0.0293, -0.0159, -0.0258],\n",
+      "        [ 0.0235,  0.0580, -0.0371,  ..., -0.0377, -0.0346,  0.0627]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.22.layer.1.EncDecAttention.v.lora_B.weight': tensor([[ 0.0026,  0.0069,  0.0095,  ..., -0.0012, -0.0045, -0.0018],\n",
+      "        [ 0.0145,  0.0194,  0.0195,  ...,  0.0182, -0.0195,  0.0177],\n",
+      "        [ 0.0069,  0.0107,  0.0172,  ...,  0.0084, -0.0127,  0.0161],\n",
+      "        ...,\n",
+      "        [ 0.0442,  0.0403, -0.0388,  ...,  0.0492, -0.0602,  0.0562],\n",
+      "        [-0.0228, -0.0222,  0.0388,  ..., -0.0255,  0.0293, -0.0299],\n",
+      "        [-0.0342, -0.0308,  0.0272,  ..., -0.0349,  0.0436, -0.0529]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.23.layer.0.SelfAttention.q.lora_A.weight': tensor([[ 0.0229,  0.0198, -0.0208,  ..., -0.0067,  0.0257, -0.0064],\n",
+      "        [-0.0030, -0.0127,  0.0058,  ..., -0.0117,  0.0089,  0.0005],\n",
+      "        [ 0.0391,  0.0234, -0.0144,  ...,  0.0433,  0.0134,  0.0189],\n",
+      "        ...,\n",
+      "        [-0.0255, -0.0331,  0.0175,  ..., -0.0176, -0.0498, -0.0518],\n",
+      "        [ 0.0129,  0.0078, -0.0142,  ..., -0.0030,  0.0345, -0.0058],\n",
+      "        [-0.0361, -0.0408,  0.0383,  ..., -0.0299, -0.0016, -0.0264]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.23.layer.0.SelfAttention.q.lora_B.weight': tensor([[-9.2302e-03,  4.7082e-04, -1.6876e-03,  ...,  1.5358e-03,\n",
+      "          3.6666e-04,  2.2397e-03],\n",
+      "        [ 6.2450e-03, -8.7633e-05,  1.2780e-03,  ..., -2.2387e-03,\n",
+      "         -7.6455e-03, -3.6240e-03],\n",
+      "        [ 7.5522e-03,  2.7108e-03,  4.8920e-03,  ..., -3.7219e-03,\n",
+      "          1.2247e-03, -5.4539e-03],\n",
+      "        ...,\n",
+      "        [ 8.3448e-03,  3.1519e-03,  2.2037e-03,  ...,  1.5882e-03,\n",
+      "          2.2505e-03, -1.0635e-03],\n",
+      "        [ 3.2382e-03,  8.2808e-03,  8.4877e-03,  ..., -5.4786e-03,\n",
+      "          9.0481e-03, -9.2424e-03],\n",
+      "        [-3.7960e-03, -1.2958e-02, -1.1950e-02,  ...,  9.1770e-03,\n",
+      "         -1.1545e-02,  7.7533e-03]], device='cuda:0'), 'base_model.model.decoder.block.23.layer.0.SelfAttention.v.lora_A.weight': tensor([[-0.0188,  0.0154,  0.0271,  ..., -0.0192,  0.0154,  0.0396],\n",
+      "        [ 0.0480,  0.0190, -0.0210,  ...,  0.0008,  0.0105, -0.0560],\n",
+      "        [-0.0263, -0.0329,  0.0212,  ..., -0.0381,  0.0282,  0.0525],\n",
+      "        ...,\n",
+      "        [-0.0505,  0.0039,  0.0447,  ..., -0.0066, -0.0126,  0.0477],\n",
+      "        [-0.0024, -0.0017,  0.0043,  ..., -0.0188, -0.0137,  0.0257],\n",
+      "        [-0.0486,  0.0265,  0.0539,  ..., -0.0064,  0.0338,  0.0425]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.23.layer.0.SelfAttention.v.lora_B.weight': tensor([[ 0.0036, -0.0072,  0.0104,  ...,  0.0092,  0.0054,  0.0088],\n",
+      "        [ 0.0053, -0.0083,  0.0139,  ...,  0.0087,  0.0064,  0.0107],\n",
+      "        [-0.0054,  0.0058, -0.0056,  ..., -0.0039, -0.0077, -0.0046],\n",
+      "        ...,\n",
+      "        [ 0.0016, -0.0062,  0.0073,  ...,  0.0109,  0.0018,  0.0105],\n",
+      "        [-0.0155,  0.0177, -0.0152,  ..., -0.0161, -0.0173, -0.0173],\n",
+      "        [ 0.0070, -0.0056, -0.0022,  ...,  0.0054,  0.0060, -0.0005]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.23.layer.1.EncDecAttention.q.lora_A.weight': tensor([[ 0.0013, -0.0205, -0.0237,  ..., -0.0034,  0.0223, -0.0525],\n",
+      "        [-0.0195, -0.0254, -0.0189,  ...,  0.0134,  0.0543, -0.0611],\n",
+      "        [ 0.0078,  0.0246, -0.0145,  ...,  0.0244,  0.0562, -0.0179],\n",
+      "        ...,\n",
+      "        [-0.0164, -0.0022, -0.0291,  ..., -0.0274, -0.0185, -0.0195],\n",
+      "        [ 0.0101,  0.0265,  0.0126,  ...,  0.0248,  0.0240, -0.0185],\n",
+      "        [-0.0226, -0.0008, -0.0094,  ...,  0.0020, -0.0359, -0.0004]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.23.layer.1.EncDecAttention.q.lora_B.weight': tensor([[-1.3781e-02, -1.4603e-02, -1.4217e-02,  ..., -9.6741e-03,\n",
+      "         -1.7460e-02,  1.1197e-02],\n",
+      "        [ 2.7672e-03,  5.4598e-03, -8.9388e-05,  ...,  2.6880e-03,\n",
+      "          1.1293e-02, -9.0429e-03],\n",
+      "        [ 3.1708e-03, -5.6693e-04, -4.4513e-03,  ...,  1.0293e-03,\n",
+      "          1.9967e-03, -4.6050e-03],\n",
+      "        ...,\n",
+      "        [-1.9154e-02,  2.7257e-03, -2.3589e-03,  ...,  1.8899e-03,\n",
+      "          6.3871e-03, -4.6289e-03],\n",
+      "        [ 1.3242e-03,  1.7306e-03, -3.5451e-03,  ..., -3.5201e-03,\n",
+      "         -5.5426e-03,  4.4722e-03],\n",
+      "        [ 1.3347e-02, -6.4381e-03, -1.2333e-02,  ..., -5.1522e-03,\n",
+      "         -1.4272e-02,  1.4489e-02]], device='cuda:0'), 'base_model.model.decoder.block.23.layer.1.EncDecAttention.v.lora_A.weight': tensor([[-0.0109, -0.0416,  0.0268,  ...,  0.0437, -0.0096, -0.0408],\n",
+      "        [-0.0006, -0.0368, -0.0175,  ...,  0.0711, -0.0297, -0.0360],\n",
+      "        [-0.0097, -0.0302, -0.0101,  ...,  0.0609, -0.0160, -0.0315],\n",
+      "        ...,\n",
+      "        [-0.0348,  0.0250,  0.0112,  ..., -0.0480,  0.0094,  0.0005],\n",
+      "        [-0.0217, -0.0358, -0.0137,  ...,  0.0188, -0.0225, -0.0686],\n",
+      "        [-0.0249, -0.0334, -0.0092,  ...,  0.0849, -0.0285, -0.0345]],\n",
+      "       device='cuda:0'), 'base_model.model.decoder.block.23.layer.1.EncDecAttention.v.lora_B.weight': tensor([[-0.0244, -0.0160, -0.0185,  ..., -0.0134, -0.0150, -0.0029],\n",
+      "        [-0.0028,  0.0118,  0.0127,  ...,  0.0023,  0.0139,  0.0067],\n",
+      "        [ 0.0218,  0.0264,  0.0261,  ...,  0.0093,  0.0255,  0.0199],\n",
+      "        ...,\n",
+      "        [-0.0049, -0.0083, -0.0071,  ..., -0.0327, -0.0097,  0.0223],\n",
+      "        [ 0.0283,  0.0257,  0.0250,  ...,  0.0135,  0.0239, -0.0075],\n",
+      "        [ 0.0004, -0.0003, -0.0002,  ..., -0.0289, -0.0049,  0.0312]],\n",
+      "       device='cuda:0')}\n"
+     ]
+    }
+   ],
+   "source": [
+    "# saving model\n",
+    "state_dict = get_pet_model_state_dict(model)\n",
+    "torch.save(state_dict, checkpoint_name)\n",
+    "print(state_dict)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "bd20cd4c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "19M\tfinancial_sentiment_analysis_lora_v1.pt\r\n"
+     ]
+    }
+   ],
+   "source": [
+    "!du -h $checkpoint_name"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "76c2fc29",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.10.5 64-bit",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.5"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/examples/pet_lora_seq2seq_accelerate_fsdp.py b/examples/pet_lora_seq2seq_accelerate_fsdp.py
new file mode 100644
index 0000000..5e0940f
--- /dev/null
+++ b/examples/pet_lora_seq2seq_accelerate_fsdp.py
@@ -0,0 +1,137 @@
+import torch
+from accelerate import Accelerator
+from torch.utils.data import DataLoader
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, default_data_collator, get_linear_schedule_with_warmup
+
+from datasets import load_dataset
+from pet import get_pet_config, get_pet_model, get_pet_model_state_dict
+from pet.utils.other import fsdp_auto_wrap_policy
+from tqdm import tqdm
+
+
+def main():
+    accelerator = Accelerator()
+    model_name_or_path = "bigscience/mt0-xxl"
+    batch_size = 16
+    text_column = "sentence"
+    label_column = "text_label"
+    max_length = 64
+    lr = 1e-3
+    num_epochs = 1
+
+    config = {"pet_type": "LORA", "task_type": "SEQ_2_SEQ_LM", "r": 8, "lora_alpha": 32, "lora_dropout": 0.1}
+    pet_config = get_pet_config(config)
+    checkpoint_name = "financial_sentiment_analysis_lora_fsdp_v1.pt"
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
+    model = get_pet_model(model, pet_config)
+    accelerator.print(model.print_trainable_parameters())
+
+    dataset = load_dataset("financial_phrasebank", "sentences_allagree")
+    dataset = dataset["train"].train_test_split(test_size=0.1)
+    dataset["validation"] = dataset["test"]
+    del dataset["test"]
+
+    classes = dataset["train"].features["label"].names
+    dataset = dataset.map(
+        lambda x: {"text_label": [classes[label] for label in x["label"]]},
+        batched=True,
+        num_proc=1,
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+
+    def preprocess_function(examples):
+        inputs = examples[text_column]
+        targets = examples[label_column]
+        model_inputs = tokenizer(
+            inputs, max_length=max_length, padding="max_length", truncation=True, return_tensors="pt"
+        )
+        labels = tokenizer(targets, max_length=3, padding="max_length", truncation=True, return_tensors="pt")
+        labels = labels["input_ids"]
+        labels[labels == tokenizer.pad_token_id] = -100
+        model_inputs["labels"] = labels
+        return model_inputs
+
+    processed_datasets = dataset.map(
+        preprocess_function,
+        batched=True,
+        num_proc=1,
+        remove_columns=dataset["train"].column_names,
+        load_from_cache_file=False,
+        desc="Running tokenizer on dataset",
+    )
+
+    train_dataset = processed_datasets["train"]
+    eval_dataset = processed_datasets["validation"]
+
+    train_dataloader = DataLoader(
+        train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True
+    )
+    eval_dataloader = DataLoader(
+        eval_dataset, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True
+    )
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=0,
+        num_training_steps=(len(train_dataloader) * num_epochs),
+    )
+
+    if accelerator.state.fsdp_plugin is not None:
+        accelerator.state.fsdp_plugin.auto_wrap_policy = fsdp_auto_wrap_policy(model)
+
+    model, train_dataloader, eval_dataloader, optimizer, lr_scheduler = accelerator.prepare(
+        model, train_dataloader, eval_dataloader, optimizer, lr_scheduler
+    )
+    accelerator.print(model)
+
+    for epoch in range(num_epochs):
+        model.train()
+        total_loss = 0
+        for step, batch in enumerate(tqdm(train_dataloader)):
+            outputs = model(**batch)
+            loss = outputs.loss
+            total_loss += loss.detach().float()
+            loss.backward()
+            optimizer.step()
+            lr_scheduler.step()
+            optimizer.zero_grad()
+
+        model.eval()
+        eval_loss = 0
+        eval_preds = []
+        for step, batch in enumerate(tqdm(eval_dataloader)):
+            with torch.no_grad():
+                outputs = model(**batch)
+            loss = outputs.loss
+            eval_loss += loss.detach().float()
+            eval_preds.extend(
+                tokenizer.batch_decode(
+                    accelerator.gather_for_metrics(torch.argmax(outputs.logits, -1)).detach().cpu().numpy(),
+                    skip_special_tokens=True,
+                )
+            )
+
+        eval_epoch_loss = eval_loss / len(train_dataloader)
+        eval_ppl = torch.exp(eval_epoch_loss)
+        train_epoch_loss = total_loss / len(eval_dataloader)
+        train_ppl = torch.exp(train_epoch_loss)
+        accelerator.print(f"{epoch=}: {train_ppl=} {train_epoch_loss=} {eval_ppl=} {eval_epoch_loss=}")
+
+        correct = 0
+        total = 0
+        for pred, true in zip(eval_preds, dataset["validation"][label_column]):
+            if pred.strip() == true.strip():
+                correct += 1
+            total += 1
+        accuracy = correct / total * 100
+        accelerator.print(f"{accuracy=}")
+        accelerator.print(f"{eval_preds[:10]=}")
+        accelerator.wait_for_everyone()
+        accelerator.save(get_pet_model_state_dict(model), checkpoint_name)
+        accelerator.wait_for_everyone()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/pet_prefix_tuning_seq2seq.ipynb b/examples/pet_prefix_tuning_seq2seq.ipynb
new file mode 100644
index 0000000..a1d5dc4
--- /dev/null
+++ b/examples/pet_prefix_tuning_seq2seq.ipynb
@@ -0,0 +1,1963 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "5f93b7d1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import AutoModelForSeq2SeqLM\n",
+    "from pet import get_pet_config,get_pet_model, get_pet_model_state_dict\n",
+    "import torch\n",
+    "from datasets import load_dataset\n",
+    "import os\n",
+    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\"\n",
+    "from transformers import AutoTokenizer\n",
+    "from torch.utils.data import DataLoader\n",
+    "from transformers import default_data_collator,get_linear_schedule_with_warmup\n",
+    "from tqdm import tqdm\n",
+    "from datasets import load_dataset\n",
+    "\n",
+    "device = \"cuda\"\n",
+    "model_name_or_path = \"t5-large\"\n",
+    "tokenizer_name_or_path = \"t5-large\"\n",
+    "\n",
+    "config = {\n",
+    "    \"pet_type\":\"PREFIX_TUNING\",\n",
+    "    \"task_type\":\"SEQ_2_SEQ_LM\",\n",
+    "    \"num_virtual_tokens\": 20\n",
+    "}\n",
+    "checkpoint_name = \"financial_sentiment_analysis_prefix_tuning_v1.pt\"\n",
+    "text_column = \"sentence\"\n",
+    "label_column = \"text_label\"\n",
+    "max_length=128\n",
+    "lr = 1e-2\n",
+    "num_epochs = 5\n",
+    "batch_size=8\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "8d0850ac",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3d73e8a9f76d4a2e918238f8ad0b2a30",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Downloading:   0%|          | 0.00/1.20k [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2965839a8e7f48fdafe2d7a909666b25",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Downloading:   0%|          | 0.00/2.95G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "trainable params: 983040 || all params: 738651136 || trainable%: 0.13308583065659835\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "PETModelForSeq2SeqLM(\n",
+       "  (base_model): T5ForConditionalGeneration(\n",
+       "    (shared): Embedding(32128, 1024)\n",
+       "    (encoder): T5Stack(\n",
+       "      (embed_tokens): Embedding(32128, 1024)\n",
+       "      (block): ModuleList(\n",
+       "        (0): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (relative_attention_bias): Embedding(32, 16)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (1): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (2): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (3): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (4): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (5): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (6): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (7): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (8): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (9): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (10): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (11): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (12): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (13): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (14): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (15): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (16): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (17): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (18): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (19): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (20): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (21): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (22): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (23): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "      )\n",
+       "      (final_layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "      (dropout): Dropout(p=0.1, inplace=False)\n",
+       "    )\n",
+       "    (decoder): T5Stack(\n",
+       "      (embed_tokens): Embedding(32128, 1024)\n",
+       "      (block): ModuleList(\n",
+       "        (0): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (relative_attention_bias): Embedding(32, 16)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (1): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (2): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (3): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (4): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (5): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (6): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (7): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (8): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (9): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (10): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (11): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (12): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (13): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (14): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (15): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (16): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (17): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (18): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (19): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (20): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (21): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (22): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "        (23): T5Block(\n",
+       "          (layer): ModuleList(\n",
+       "            (0): T5LayerSelfAttention(\n",
+       "              (SelfAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (1): T5LayerCrossAttention(\n",
+       "              (EncDecAttention): T5Attention(\n",
+       "                (q): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (k): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (v): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "                (o): Linear(in_features=1024, out_features=1024, bias=False)\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (2): T5LayerFF(\n",
+       "              (DenseReluDense): T5DenseActDense(\n",
+       "                (wi): Linear(in_features=1024, out_features=4096, bias=False)\n",
+       "                (wo): Linear(in_features=4096, out_features=1024, bias=False)\n",
+       "                (dropout): Dropout(p=0.1, inplace=False)\n",
+       "                (act): ReLU()\n",
+       "              )\n",
+       "              (layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "        )\n",
+       "      )\n",
+       "      (final_layer_norm): FusedRMSNorm(torch.Size([1024]), eps=1e-06, elementwise_affine=True)\n",
+       "      (dropout): Dropout(p=0.1, inplace=False)\n",
+       "    )\n",
+       "    (lm_head): Linear(in_features=1024, out_features=32128, bias=False)\n",
+       "  )\n",
+       "  (word_embeddings): Embedding(32128, 1024)\n",
+       "  (prompt_encoder): PrefixEncoder(\n",
+       "    (embedding): Embedding(20, 49152)\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# creating model\n",
+    "pet_config = get_pet_config(config)\n",
+    "\n",
+    "model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)\n",
+    "model = get_pet_model(model, pet_config)\n",
+    "model.print_trainable_parameters()\n",
+    "model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "4ee2babf",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/sourab/miniconda3/envs/ml/lib/python3.10/site-packages/huggingface_hub/utils/_deprecation.py:97: FutureWarning: Deprecated argument(s) used in 'dataset_info': token. Will not be supported from version '0.12'.\n",
+      "  warnings.warn(message, FutureWarning)\n",
+      "Found cached dataset financial_phrasebank (/home/sourab/.cache/huggingface/datasets/financial_phrasebank/sentences_allagree/1.0.0/550bde12e6c30e2674da973a55f57edde5181d53f5a5a34c1531c53f93b7e141)\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e3f8b8faca0a4112b2c3499faee9544b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/1 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "935c8aebde284a5784348588e0bb013a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/3 [00:00<?, ?ba/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e3487cd55f6847588492bf7fa51348ca",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/1 [00:00<?, ?ba/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "{'sentence': 'ADPnews - Feb 5 , 2010 - Finnish real estate investor Sponda Oyj HEL : SDA1V said today that it slipped to a net loss of EUR 81.5 million USD 11.8 m in 2009 from a profit of EUR 29.3 million in 2008 .',\n",
+       " 'label': 0,\n",
+       " 'text_label': 'negative'}"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# loading dataset\n",
+    "dataset = load_dataset(\"financial_phrasebank\", 'sentences_allagree')\n",
+    "dataset = dataset[\"train\"].train_test_split(test_size=0.1)\n",
+    "dataset[\"validation\"] = dataset[\"test\"]\n",
+    "del(dataset[\"test\"])\n",
+    "\n",
+    "classes = dataset[\"train\"].features[\"label\"].names\n",
+    "dataset = dataset.map(\n",
+    "    lambda x: {\"text_label\": [classes[label] for label in x[\"label\"]]},\n",
+    "    batched=True,\n",
+    "    num_proc=1,\n",
+    "    \n",
+    ")\n",
+    "\n",
+    "dataset[\"train\"][0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "adf9608c",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2ce088f4437d4e2c80c267332a5b84e5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Downloading:   0%|          | 0.00/792k [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4e5f69b61f194220b39336e48edd2f9e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Downloading:   0%|          | 0.00/1.39M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/sourab/transformers/src/transformers/models/t5/tokenization_t5_fast.py:156: FutureWarning: This tokenizer was incorrectly instantiated with a model max length of 512 which will be corrected in Transformers v5.\n",
+      "For now, this behavior is kept to avoid breaking backwards compatibility when padding/encoding with `truncation is True`.\n",
+      "- Be aware that you SHOULD NOT rely on t5-large automatically truncating your input to 512 when padding/encoding.\n",
+      "- If you want to encode/pad to sequences longer than 512 you can either instantiate this tokenizer with `model_max_length` or pass `max_length` when encoding/padding.\n",
+      "- To avoid this warning, please instantiate this tokenizer with `model_max_length` set to your preferred value.\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "230c5631891e4ea8ac7a1b39f315a4f0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Running tokenizer on dataset:   0%|          | 0/3 [00:00<?, ?ba/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b581e5677d2a45459ceb725534ed0891",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Running tokenizer on dataset:   0%|          | 0/1 [00:00<?, ?ba/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# data preprocessing\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n",
+    "def preprocess_function(examples):\n",
+    "    inputs = examples[text_column]\n",
+    "    targets = examples[label_column]\n",
+    "    model_inputs = tokenizer(inputs, max_length=max_length, padding=\"max_length\", truncation=True, return_tensors=\"pt\")\n",
+    "    labels = tokenizer(targets, max_length=2, padding=\"max_length\", truncation=True, return_tensors=\"pt\")\n",
+    "    labels = labels[\"input_ids\"]\n",
+    "    labels[labels==tokenizer.pad_token_id] = -100\n",
+    "    model_inputs[\"labels\"] = labels\n",
+    "    return model_inputs\n",
+    "\n",
+    "processed_datasets = dataset.map(\n",
+    "            preprocess_function,\n",
+    "            batched=True,\n",
+    "            num_proc=1,\n",
+    "            remove_columns=dataset[\"train\"].column_names,\n",
+    "            load_from_cache_file=False,\n",
+    "            desc=\"Running tokenizer on dataset\",\n",
+    "        )\n",
+    "\n",
+    "train_dataset = processed_datasets[\"train\"]\n",
+    "eval_dataset = processed_datasets[\"validation\"]\n",
+    "\n",
+    "train_dataloader = DataLoader(\n",
+    "        train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True\n",
+    "    )\n",
+    "eval_dataloader = DataLoader(eval_dataset, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True)\n",
+    "\n",
+    "\n",
+    "\n",
+    "    "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "f733a3c6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# optimizer and lr scheduler\n",
+    "optimizer = torch.optim.AdamW(model.parameters(), lr=lr)\n",
+    "lr_scheduler = get_linear_schedule_with_warmup(\n",
+    "    optimizer=optimizer,\n",
+    "    num_warmup_steps=0,\n",
+    "    num_training_steps=(len(train_dataloader) * num_epochs),\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "6b3a4090",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████████████████████████| 255/255 [00:20<00:00, 12.27it/s]\n",
+      "100%|███████████████████████████████████████████████████████████████| 29/29 [00:01<00:00, 17.32it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=0: train_ppl=tensor(2697769., device='cuda:0') train_epoch_loss=tensor(14.8079, device='cuda:0') eval_ppl=tensor(1.0089, device='cuda:0') eval_epoch_loss=tensor(0.0089, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████████████████████████| 255/255 [00:19<00:00, 12.75it/s]\n",
+      "100%|███████████████████████████████████████████████████████████████| 29/29 [00:01<00:00, 17.33it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=1: train_ppl=tensor(2.9475, device='cuda:0') train_epoch_loss=tensor(1.0809, device='cuda:0') eval_ppl=tensor(1.0072, device='cuda:0') eval_epoch_loss=tensor(0.0072, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████████████████████████| 255/255 [00:20<00:00, 12.71it/s]\n",
+      "100%|███████████████████████████████████████████████████████████████| 29/29 [00:01<00:00, 17.31it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=2: train_ppl=tensor(2.0588, device='cuda:0') train_epoch_loss=tensor(0.7221, device='cuda:0') eval_ppl=tensor(1.0055, device='cuda:0') eval_epoch_loss=tensor(0.0054, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████████████████████████| 255/255 [00:20<00:00, 12.70it/s]\n",
+      "100%|███████████████████████████████████████████████████████████████| 29/29 [00:01<00:00, 17.32it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=3: train_ppl=tensor(1.7939, device='cuda:0') train_epoch_loss=tensor(0.5844, device='cuda:0') eval_ppl=tensor(1.0063, device='cuda:0') eval_epoch_loss=tensor(0.0063, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████████████████████████| 255/255 [00:19<00:00, 13.01it/s]\n",
+      "100%|███████████████████████████████████████████████████████████████| 29/29 [00:01<00:00, 17.33it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=4: train_ppl=tensor(1.7740, device='cuda:0') train_epoch_loss=tensor(0.5732, device='cuda:0') eval_ppl=tensor(1.0062, device='cuda:0') eval_epoch_loss=tensor(0.0061, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "# training and evaluation\n",
+    "model = model.to(device)\n",
+    "\n",
+    "for epoch in range(num_epochs):\n",
+    "    model.train()\n",
+    "    total_loss = 0\n",
+    "    for step, batch in enumerate(tqdm(train_dataloader)):\n",
+    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
+    "        outputs = model(**batch)\n",
+    "        loss = outputs.loss\n",
+    "        total_loss += loss.detach().float()\n",
+    "        loss.backward()\n",
+    "        optimizer.step()\n",
+    "        lr_scheduler.step()\n",
+    "        optimizer.zero_grad()\n",
+    "\n",
+    "    model.eval()\n",
+    "    eval_loss = 0\n",
+    "    eval_preds = []\n",
+    "    for step, batch in enumerate(tqdm(eval_dataloader)):\n",
+    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
+    "        with torch.no_grad():\n",
+    "            outputs = model(**batch)\n",
+    "        loss = outputs.loss\n",
+    "        eval_loss += loss.detach().float()\n",
+    "        eval_preds.extend(tokenizer.batch_decode(torch.argmax(outputs.logits, -1).detach().cpu().numpy(), skip_special_tokens=True))\n",
+    "\n",
+    "    eval_epoch_loss = eval_loss/len(train_dataloader)\n",
+    "    eval_ppl = torch.exp(eval_epoch_loss)\n",
+    "    train_epoch_loss = total_loss/len(eval_dataloader)\n",
+    "    train_ppl = torch.exp(train_epoch_loss)\n",
+    "    print(f\"{epoch=}: {train_ppl=} {train_epoch_loss=} {eval_ppl=} {eval_epoch_loss=}\")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "6cafa67b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "accuracy=96.47577092511013 % on the evaluation dataset\n",
+      "eval_preds[:10]=['neutral', 'neutral', 'neutral', 'negative', 'neutral', 'neutral', 'neutral', 'neutral', 'positive', 'positive']\n",
+      "dataset['validation']['text_label'][:10]=['neutral', 'neutral', 'neutral', 'negative', 'neutral', 'neutral', 'neutral', 'neutral', 'positive', 'positive']\n"
+     ]
+    }
+   ],
+   "source": [
+    "# print accuracy\n",
+    "correct =0\n",
+    "total = 0\n",
+    "for pred,true in zip(eval_preds, dataset[\"validation\"][\"text_label\"]):\n",
+    "    if pred.strip()==true.strip():\n",
+    "        correct+=1\n",
+    "    total+=1  \n",
+    "accuracy = correct/total*100\n",
+    "print(f\"{accuracy=} % on the evaluation dataset\")\n",
+    "print(f\"{eval_preds[:10]=}\")\n",
+    "print(f\"{dataset['validation']['text_label'][:10]=}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "a8de6005",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'prompt_embeddings': tensor([[-0.3165, -0.8389,  0.3262,  ..., -1.5049, -1.6963,  0.3444],\n",
+      "        [-1.8359,  1.1936,  1.0483,  ...,  0.6197, -0.4452,  0.5844],\n",
+      "        [-0.6027,  0.3246, -1.5601,  ..., -0.3645,  0.2329,  0.3402],\n",
+      "        ...,\n",
+      "        [-1.9525, -0.5035,  0.8474,  ...,  0.4793, -0.0789, -0.9305],\n",
+      "        [-1.9741,  0.5242, -2.0594,  ..., -0.7970, -0.4889,  2.7323],\n",
+      "        [ 0.9355, -0.2714,  0.4610,  ...,  0.2692, -1.5801, -1.6405]])}\n"
+     ]
+    }
+   ],
+   "source": [
+    "# saving model\n",
+    "state_dict = get_pet_model_state_dict(model)\n",
+    "torch.save(state_dict, checkpoint_name)\n",
+    "print(state_dict)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "bd20cd4c",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3,8M\tfinancial_sentiment_analysis_prefix_tuning_v1.pt\r\n"
+     ]
+    }
+   ],
+   "source": [
+    "!du -h $checkpoint_name"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "76c2fc29",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/setup.py b/setup.py
index 6f74e46..fcd1b39 100644
--- a/setup.py
+++ b/setup.py
@@ -39,9 +39,10 @@ setup(
         "packaging>=20.0",
         "psutil",
         "pyyaml",
-        "torch>=1.4.0",
+        "torch>=1.13.0",
         "transformers",
         "accelerate",
+        "loralib",
     ],
     extras_require=extras,
     classifiers=[
diff --git a/src/pet/mapping.py b/src/pet/mapping.py
index 0cb2210..115372e 100644
--- a/src/pet/mapping.py
+++ b/src/pet/mapping.py
@@ -18,6 +18,7 @@ PET_TYPE_TO_CONFIG_MAPPING = {
 
 TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
     "t5": ["q", "v"],
+    "mt5": ["q", "v"],
     "bart": ["q_proj", "v_proj"],
     "gpt2": ["c_attn"],
     "bloom": ["query_key_value"],
@@ -27,6 +28,7 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
     "gpt_neo": ["q_proj", "v_proj"],
     "bert": ["query", "value"],
     "roberta": ["query", "value"],
+    "xlm-roberta": ["query", "value"],
     "electra": ["query", "value"],
     "deberta-v2": ["query_proj", "value_proj"],
     "deberta": ["in_proj"],
@@ -34,6 +36,13 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
 
 
 def get_pet_config(config_dict):
+    """
+    Returns a PET config object from a dictionary.
+
+    Args:
+        config_dict (:obj:`Dict[str, Any]`):
+    """
+
     return PET_TYPE_TO_CONFIG_MAPPING[config_dict["pet_type"]](**config_dict)
 
 
@@ -73,8 +82,8 @@ def _prepare_prompt_learning_config(pet_config, model_config):
             raise ValueError("Please specify `num_attention_heads` in `pet_config`")
         pet_config.num_attention_heads = num_attention_heads
 
-    if pet_config.encoder_hidden_size is None:
-        pet_config.encoder_hidden_size = token_dim
+    if getattr(pet_config, "encoder_hidden_size", None) is None:
+        setattr(pet_config, "encoder_hidden_size", token_dim)
 
     return pet_config
 
@@ -93,6 +102,14 @@ def _prepare_lora_config(pet_config, model_config):
 
 
 def get_pet_model(model, pet_config):
+    """
+    Returns a PET model object from a model and a config.
+
+    Args:
+        model (:obj:`transformers.PreTrainedModel`):
+        pet_config (:obj:`transformers.PETConfig`):
+    """
+
     model_config = model.config.to_dict()
     if pet_config.pet_type != PETType.LORA:
         pet_config = _prepare_prompt_learning_config(pet_config, model_config)
diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py
index 1398146..dbe2b6e 100644
--- a/src/pet/pet_model.py
+++ b/src/pet/pet_model.py
@@ -7,28 +7,50 @@ from transformers import PreTrainedModel
 from transformers.modeling_outputs import SequenceClassifierOutput
 
 from .tuners import LoRAModel, PrefixEncoder, PromptEmbedding, PromptEncoder
-from .utils import PETConfig, PETType, TaskType, shift_tokens_right
+from .utils import PETConfig, PETType, TaskType, _set_trainable, shift_tokens_right
 
 
 class PETModel(torch.nn.Module):
+    """
+    Parameter Efficient Tuning Model. Base model encompassing various PET methods.
+
+    Args:
+        model (:obj:`PreTrainedModel`): The base transformer model used for PET.
+        pet_config (:obj:`PETConfig`): The configuration of the PET model.
+
+
+    Attributes:
+        base_model (:obj:`PreTrainedModel`): The base transformer model used for PET. pet_config (:obj:`PETConfig`):
+        The configuration of the PET model. modules_to_save (:obj:`list` of :obj:`str`): The list of sub-module names
+        to save when saving the model. prompt_encoder (:obj:`PromptEncoder`): The prompt encoder used for PET if
+        `pet_config.pet_type != PETType.LORA`. prompt_tokens (:obj:`torch.Tensor`): The virtual prompt tokens used for
+        PET if `pet_config.pet_type != PETType.LORA`. transformer_backbone_name (:obj:`str`): The name of the
+        transformer backbone in the base model
+            if `pet_config.pet_type != PETType.LORA`.
+        word_embeddings (:obj:`torch.nn.Embedding`): The word embeddings of the transformer backbone
+            in the base model if `pet_config.pet_type != PETType.LORA`.
+    """
+
     def __init__(self, model, pet_config: PETConfig):
         super().__init__()
         self.pet_config = pet_config
         self.base_model = model
+        self.config = self.base_model.config
         self.modules_to_save = None
         if pet_config.pet_type != PETType.LORA:
             self._setup_prompt_encoder()
         else:
             self.base_model = LoRAModel(pet_config, model)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     def _setup_prompt_encoder(self):
         num_transformer_submodules = 0
         transformer_backbone = None
         for name, module in self.base_model.named_children():
+            for param in module.parameters():
+                param.requires_grad = False
             if isinstance(module, PreTrainedModel):
                 # Make sure to freeze Tranformers model
-                for param in module.parameters():
-                    param.requires_grad = False
                 if transformer_backbone is None:
                     transformer_backbone = module
                     self.transformer_backbone_name = name
@@ -54,14 +76,21 @@ class PETModel(torch.nn.Module):
         ).long()
 
     def get_prompt_embedding_to_save(self):
-        prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.base_model.device)
+        """
+        Returns the prompt embedding to save when saving the model. Only applocable when `pet_config.pet_type !=
+        PETType.LORA`.
+        """
+        prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.device)
         if self.pet_config.pet_type == PETType.PREFIX_TUNING:
             prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens]
         prompt_embeddings = self.prompt_encoder(prompt_tokens)
         return prompt_embeddings[0].detach().cpu()
 
     def get_prompt(self, batch_size):
-        prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.base_model.device)
+        """
+        Returns the virtual prompts to use for PET. Only applocable when `pet_config.pet_type != PETType.LORA`.
+        """
+        prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.device)
         if self.pet_config.pet_type == PETType.PREFIX_TUNING:
             prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens]
             if self.pet_config.inference_mode:
@@ -92,6 +121,9 @@ class PETModel(torch.nn.Module):
             return prompts
 
     def print_trainable_parameters(self):
+        """
+        Prints the number of trainable parameters in the model.
+        """
         trainable_params = 0
         all_param = 0
         for _, param in self.named_parameters():
@@ -102,11 +134,42 @@ class PETModel(torch.nn.Module):
             f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
         )
 
+    def __getattr__(self, name: str):
+        """Forward missing attributes to the wrapped module."""
+        try:
+            return super().__getattr__(name)  # defer to nn.Module's logic
+        except AttributeError:
+            return getattr(self.base_model, name)
+
 
 class PETModelForSequenceClassification(PETModel):
+    """
+    PET model for sequence classification tasks.
+
+    Args:
+        model (:obj:`PreTrainedModel`): Base transformer model
+        pet_config (:obj:`PETConfig`): PET config.
+
+    Attributes:
+        config (:obj:`PretrainedConfig`): The configuration object of the base model. cls_layer_name (:obj:`str`): The
+        name of the classification layer.
+
+    Example::
+
+        >>> from transformers import AutoModelForSequenceClassification >>> from pet import
+        PETModelForSequenceClassification, get_pet_config >>> config = {
+                'pet_type': 'PREFIX_TUNING', 'task_type': 'SEQ_CLS', 'inference_mode': False, 'num_virtual_tokens': 20,
+                'token_dim': 768, 'num_transformer_submodules': 1, 'num_attention_heads': 12, 'num_layers': 12,
+                'encoder_hidden_size': 768, 'prefix_projection': False, 'postprocess_past_key_value_function': None
+            }
+        >>> pet_config = get_pet_config(config) >>> model =
+        AutoModelForSequenceClassification.from_pretrained("bert-base-cased") >>> pet_model =
+        PETModelForSequenceClassification(model, pet_config) >>> pet_model.print_trainable_parameters() trainable
+        params: 370178 || all params: 108680450 || trainable%: 0.3406113979101117
+    """
+
     def __init__(self, model, pet_config: PETConfig):
         super().__init__(model, pet_config)
-        self.config = self.base_model.config
         self.modules_to_save = ["classifier"]
 
         for name, module in self.base_model.named_children():
@@ -114,6 +177,9 @@ class PETModelForSequenceClassification(PETModel):
                 self.cls_layer_name = name
                 break
 
+        # to make sure classifier layer is trainable
+        _set_trainable(self.base_model)
+
     def forward(
         self,
         input_ids=None,
@@ -142,9 +208,7 @@ class PETModelForSequenceClassification(PETModel):
         batch_size = input_ids.shape[0]
         if attention_mask is not None:
             # concat prompt attention mask
-            prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
-                self.base_model.device
-            )
+            prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(self.device)
             attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
         if kwargs.get("position_ids", None) is not None:
             warnings.warn("Position ids are not supported for parameter efficient tuning. Ignoring position ids.")
@@ -160,12 +224,12 @@ class PETModelForSequenceClassification(PETModel):
         )
 
         if self.pet_config.pet_type == PETType.PREFIX_TUNING:
-            return self.prefix_tuning_forward(input_ids=input_ids, **kwargs)
+            return self._prefix_tuning_forward(input_ids=input_ids, **kwargs)
         else:
             if kwargs.get("token_type_ids", None) is not None:
                 kwargs["token_type_ids"] = torch.cat(
                     (
-                        torch.zeros(batch_size, self.pet_config.num_virtual_tokens).to(self.base_model.device),
+                        torch.zeros(batch_size, self.pet_config.num_virtual_tokens).to(self.device),
                         kwargs["token_type_ids"],
                     ),
                     dim=1,
@@ -176,7 +240,7 @@ class PETModelForSequenceClassification(PETModel):
             inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
             return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
 
-    def prefix_tuning_forward(
+    def _prefix_tuning_forward(
         self,
         input_ids=None,
         attention_mask=None,
@@ -249,9 +313,29 @@ class PETModelForSequenceClassification(PETModel):
 
 
 class PETModelForCausalLM(PETModel):
+    """
+    PET model for Causal LM
+
+    Args:
+        model (:obj:`PreTrainedModel`): Base transformer model
+        pet_config (:obj:`PETConfig`): PET config.
+
+
+    Example::
+
+        >>> from transformers import AutoModelForCausalLM >>> from pet import PETModelForCausalLM, get_pet_config >>>
+        config = {
+                'pet_type': 'PREFIX_TUNING', 'task_type': 'CAUSAL_LM', 'inference_mode': False, 'num_virtual_tokens':
+                20, 'token_dim': 1280, 'num_transformer_submodules': 1, 'num_attention_heads': 20, 'num_layers': 36,
+                'encoder_hidden_size': 1280, 'prefix_projection': False, 'postprocess_past_key_value_function': None
+            }
+        >>> pet_config = get_pet_config(config) >>> model = AutoModelForCausalLM.from_pretrained("gpt2-large") >>>
+        pet_model = PETModelForCausalLM(model, pet_config) >>> pet_model.print_trainable_parameters() trainable params:
+        1843200 || all params: 775873280 || trainable%: 0.23756456724479544
+    """
+
     def __init__(self, model, pet_config: PETConfig):
         super().__init__(model, pet_config)
-        self.config = self.base_model.config
 
     def forward(
         self,
@@ -279,9 +363,7 @@ class PETModelForCausalLM(PETModel):
         batch_size = input_ids.shape[0]
         if attention_mask is not None:
             # concat prompt attention mask
-            prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
-                self.base_model.device
-            )
+            prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(self.device)
             attention_mask = torch.cat((prefix_attention_mask, attention_mask), dim=1)
 
         if kwargs.get("position_ids", None) is not None:
@@ -308,9 +390,7 @@ class PETModelForCausalLM(PETModel):
                 inputs_embeds = self.word_embeddings(input_ids)
             # concat prompt labels
             if labels is not None:
-                prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to(
-                    self.base_model.device
-                )
+                prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to(self.device)
                 kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1)
             prompts = self.get_prompt(batch_size=batch_size)
             inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
@@ -318,9 +398,29 @@ class PETModelForCausalLM(PETModel):
 
 
 class PETModelForSeq2SeqLM(PETModel):
+    """
+    PET model for Seq2Seq LM
+
+    Args:
+        model (:obj:`PreTrainedModel`): Base transformer model
+        pet_config (:obj:`PETConfig`): PET config.
+
+
+    Example::
+
+        >>> from transformers import AutoModelForSeq2SeqLM >>> from pet import PETModelForSeq2SeqLM, get_pet_config >>>
+        config = {
+                'pet_type': 'LORA', 'task_type': 'SEQ_2_SEQ_LM', 'inference_mode': False, 'r': 8, 'target_modules':
+                ['q', 'v'], 'lora_alpha': 32, 'lora_dropout': 0.1, 'merge_weights': False, 'fan_in_fan_out': False,
+                'enable_lora': None, 'bias': 'none'
+            }
+        >>> pet_config = get_pet_config(config) >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") >>>
+        pet_model = PETModelForSeq2SeqLM(model, pet_config) >>> pet_model.print_trainable_parameters() trainable
+        params: 884736 || all params: 223843584 || trainable%: 0.3952474242013566
+    """
+
     def __init__(self, model, pet_config: PETConfig):
         super().__init__(model, pet_config)
-        self.config = self.base_model.config
 
     def forward(
         self,
@@ -354,9 +454,7 @@ class PETModelForSeq2SeqLM(PETModel):
         batch_size = input_ids.shape[0]
         if decoder_attention_mask is not None:
             # concat prompt attention mask
-            prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
-                self.base_model.device
-            )
+            prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(self.device)
             decoder_attention_mask = torch.cat((prefix_attention_mask, decoder_attention_mask), dim=1)
 
         if kwargs.get("position_ids", None) is not None:
@@ -392,15 +490,11 @@ class PETModelForSeq2SeqLM(PETModel):
 
             if attention_mask is not None:
                 # concat prompt attention mask
-                prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(
-                    self.base_model.device
-                )
+                prefix_attention_mask = torch.ones(batch_size, self.pet_config.num_virtual_tokens).to(self.device)
                 kwargs["attention_mask"] = torch.cat((prefix_attention_mask, attention_mask), dim=1)
             # concat prompt labels
             if labels is not None:
-                prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to(
-                    self.base_model.device
-                )
+                prefix_labels = torch.full((batch_size, self.pet_config.num_virtual_tokens), -100).to(self.device)
                 kwargs["labels"] = torch.cat((prefix_labels, labels), dim=1)
             prompts = self.get_prompt(batch_size=batch_size)
             inputs_embeds = torch.cat((prompts[:, : self.pet_config.num_virtual_tokens], inputs_embeds), dim=1)
diff --git a/src/pet/tuners/lora.py b/src/pet/tuners/lora.py
index eb30bad..cf57404 100644
--- a/src/pet/tuners/lora.py
+++ b/src/pet/tuners/lora.py
@@ -1,11 +1,14 @@
 # todo
+import math
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import List, Optional
 
 import torch
+import torch.nn as nn
+import torch.nn.functional as F
 from transformers.pytorch_utils import Conv1D
 
-import loralib as lora
+import loralib as lora  # noqa: F401
 from loralib import mark_only_lora_as_trainable
 
 from ..utils import PETConfig
@@ -13,6 +16,21 @@ from ..utils import PETConfig
 
 @dataclass
 class LoRAConfig(PETConfig):
+    """
+    This is the configuration class to store the configuration of a :class:`~pet.LoRA`.
+
+    Args:
+        r: (:obj:`init`): LoRA attention dimension
+        target_modules (:obj: list of :obj: str): The names of the modules to apply LoRA to.
+        lora_alpha (:obj: float): The alpha parameter for LoRA scaling.
+        lora_dropout (:obj: float): The dropout probability for LoRA layers.
+        merge_weights (:obj: bool):
+            Whether to merge the weights of the LoRA layers with the base transformer model in `eval` mode.
+        fan_in_fan_out (:obj: bool): Set this to True if the layer to replace stores weight like (fan_in, fan_out)
+        enable_lora (:obj: list of :obj: bool): Used with `lora.MergedLinear`.
+        bias (:obj: str): Bias type for LoRA. Can be 'none', 'all' or 'lora_only'
+    """
+
     r: int = field(default=8, metadata={"help": "LoRA attention dimension"})
     target_modules: Optional[list] = field(default=None, metadata={"help": "List of modules to replace with LoRA"})
     lora_alpha: int = field(default=None, metadata={"help": "LoRA alpha"})
@@ -29,14 +47,37 @@ class LoRAConfig(PETConfig):
 
 
 class LoRAModel(torch.nn.Module):
+    """
+    Creates Low Rank Adapter (LoRA) model from a pretrained transformers model.
+
+    Args:
+        model (:obj:`transformers.PreTrainedModel`): The model to be adapted.
+        config (:obj:`LoRAConfig`): The configuration of the LoRA model.
+
+    Returns:
+        :obj:`torch.nn.Module`: The LoRA model.
+
+    Example::
+
+        >>> from transformers import AutoModelForSeq2SeqLM, LoRAConfig >>> from pet import LoRAModel, LoRAConfig >>>
+        config = LoRAConfig(
+            pet_type="LORA", task_type="SEQ_2_SEQ_LM", r=8, lora_alpha=32, target_modules=["q", "v"],
+            lora_dropout=0.01, )
+        >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base") >>> lora_model = LoRAModel(config, model)
+
+    Attributes:
+        model (:obj:`transformers.PreTrainedModel`): The model to be adapted. config (:obj:`LoRAConfig`): The
+        configuration of the LoRA model.
+    """
+
     def __init__(self, config, model):
         super().__init__()
         self.config = config
         self.model = model
-        self.find_and_replace()
+        self._find_and_replace()
         mark_only_lora_as_trainable(self.model, self.config.bias)
 
-    def find_and_replace(self):
+    def _find_and_replace(self):
         kwargs = {
             "r": self.config.r,
             "lora_alpha": self.config.lora_alpha,
@@ -47,23 +88,23 @@ class LoRAModel(torch.nn.Module):
         key_list = [key for key, _ in self.model.named_modules()]
         for key in key_list:
             if any(key.endswith(target_key) for target_key in self.config.target_modules):
-                parent, target, target_name = self.get_submodules(key)
-                # print(parent, target, target_name)
+                parent, target, target_name = self._get_submodules(key)
+                bias = target.bias is not None
                 if isinstance(target, torch.nn.Linear):
-                    new_module = lora.Linear(target.in_features, target.out_features, **kwargs)
+                    new_module = Linear(target.in_features, target.out_features, bias=bias, **kwargs)
                 elif isinstance(target, Conv1D):
                     kwargs.update({"enable_lora": self.config.enable_lora})
                     in_features, out_features = target.weight.shape
-                    new_module = lora.MergedLinear(in_features, out_features, **kwargs)
-                self.replace_module(parent, target_name, new_module, target)
+                    new_module = MergedLinear(in_features, out_features, bias=bias, **kwargs)
+                self._replace_module(parent, target_name, new_module, target)
 
-    def get_submodules(self, key):
+    def _get_submodules(self, key):
         parent = self.model.get_submodule(".".join(key.split(".")[:-1]))
         target_name = key.split(".")[-1]
         target = self.model.get_submodule(key)
         return parent, target, target_name
 
-    def replace_module(self, parent_module, child_name, new_module, old_module):
+    def _replace_module(self, parent_module, child_name, new_module, old_module):
         setattr(parent_module, child_name, new_module)
         new_module.weight = old_module.weight
         if old_module.bias is not None:
@@ -71,3 +112,215 @@ class LoRAModel(torch.nn.Module):
 
     def forward(self, *args, **kwargs):
         return self.model(*args, **kwargs)
+
+    def __getattr__(self, name: str):
+        """Forward missing attributes to the wrapped module."""
+        try:
+            return super().__getattr__(name)  # defer to nn.Module's logic
+        except AttributeError:
+            return getattr(self.model, name)
+
+
+# Below code is copied from https://github.com/microsoft/LoRA/blob/main/loralib/layers.py
+# and modified to work with PyTorch FSDP
+
+#  ------------------------------------------------------------------------------------------
+#  Copyright (c) Microsoft Corporation. All rights reserved.
+#  Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
+#  ------------------------------------------------------------------------------------------
+class LoRALayer:
+    def __init__(
+        self,
+        r: int,
+        lora_alpha: int,
+        lora_dropout: float,
+        merge_weights: bool,
+    ):
+        self.r = r
+        self.lora_alpha = lora_alpha
+        # Optional dropout
+        if lora_dropout > 0.0:
+            self.lora_dropout = nn.Dropout(p=lora_dropout)
+        else:
+            self.lora_dropout = lambda x: x
+        # Mark the weight as unmerged
+        self.merged = False
+        self.merge_weights = merge_weights
+
+
+class Linear(nn.Linear, LoRALayer):
+    # LoRA implemented in a dense layer
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        r: int = 0,
+        lora_alpha: int = 1,
+        lora_dropout: float = 0.0,
+        fan_in_fan_out: bool = False,  # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
+        merge_weights: bool = True,
+        **kwargs,
+    ):
+        nn.Linear.__init__(self, in_features, out_features, **kwargs)
+        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights)
+
+        self.fan_in_fan_out = fan_in_fan_out
+        # Actual trainable parameters
+        if r > 0:
+            self.lora_A = nn.Linear(in_features, r, bias=False)
+            self.lora_B = nn.Linear(r, out_features, bias=False)
+            self.scaling = self.lora_alpha / self.r
+            # Freezing the pre-trained weight matrix
+            self.weight.requires_grad = False
+        self.reset_parameters()
+        if fan_in_fan_out:
+            self.weight.data = self.weight.data.T
+
+    def reset_parameters(self):
+        nn.Linear.reset_parameters(self)
+        if hasattr(self, "lora_A"):
+            # initialize A the same way as the default for nn.Linear and B to zero
+            nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
+            nn.init.zeros_(self.lora_B.weight)
+
+    def train(self, mode: bool = True):
+        def T(w):
+            return w.T if self.fan_in_fan_out else w
+
+        nn.Linear.train(self, mode)
+        self.lora_A.train(mode)
+        self.lora_B.train(mode)
+        if self.merge_weights and self.merged:
+            # Make sure that the weights are not merged
+            if self.r > 0:
+                self.weight.data -= T(self.lora_B.weight @ self.lora_A.weight) * self.scaling
+            self.merged = False
+
+    def eval(self):
+        def T(w):
+            return w.T if self.fan_in_fan_out else w
+
+        nn.Linear.eval(self)
+        self.lora_A.eval()
+        self.lora_B.eval()
+        if self.merge_weights and not self.merged:
+            # Merge the weights and mark it
+            if self.r > 0:
+                self.weight.data += T(self.lora_B.weight @ self.lora_A.weight) * self.scaling
+            self.merged = True
+
+    def forward(self, x: torch.Tensor):
+        def T(w):
+            return w.T if self.fan_in_fan_out else w
+
+        if self.r > 0 and not self.merged:
+            result = F.linear(x, T(self.weight), bias=self.bias)
+            if self.r > 0:
+                result += self.lora_B(self.lora_A(self.lora_dropout(x))) * self.scaling
+            return result
+        else:
+            return F.linear(x, T(self.weight), bias=self.bias)
+
+
+class MergedLinear(nn.Linear, LoRALayer):
+    # LoRA implemented in a dense layer
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        r: int = 0,
+        lora_alpha: int = 1,
+        lora_dropout: float = 0.0,
+        enable_lora: List[bool] = [False],
+        fan_in_fan_out: bool = False,
+        merge_weights: bool = True,
+        **kwargs,
+    ):
+        nn.Linear.__init__(self, in_features, out_features, **kwargs)
+        LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights)
+        assert out_features % len(enable_lora) == 0, "The length of enable_lora must divide out_features"
+        self.enable_lora = enable_lora
+        self.fan_in_fan_out = fan_in_fan_out
+        # Actual trainable parameters
+        if r > 0 and any(enable_lora):
+            self.lora_A = nn.Linear(in_features, r * sum(enable_lora), bias=False)
+            self.lora_B = nn.Conv1d(
+                r * sum(enable_lora),
+                out_features // len(enable_lora) * sum(enable_lora),
+                kernel_size=1,
+                groups=2,
+                bias=False,
+            )
+            self.scaling = self.lora_alpha / self.r
+            # Freezing the pre-trained weight matrix
+            self.weight.requires_grad = False
+            # Compute the indices
+            self.lora_ind = self.weight.new_zeros((out_features,), dtype=torch.bool).view(len(enable_lora), -1)
+            self.lora_ind[enable_lora, :] = True
+            self.lora_ind = self.lora_ind.view(-1)
+        self.reset_parameters()
+        if fan_in_fan_out:
+            self.weight.data = self.weight.data.T
+
+    def reset_parameters(self):
+        nn.Linear.reset_parameters(self)
+        if hasattr(self, "lora_A"):
+            # initialize A the same way as the default for nn.Linear and B to zero
+            nn.init.kaiming_uniform_(self.lora_A.weight, a=math.sqrt(5))
+            nn.init.zeros_(self.lora_B.weight)
+
+    def zero_pad(self, x):
+        result = x.new_zeros((*x.shape[:-1], self.out_features))
+        result = result.view(-1, self.out_features)
+        result[:, self.lora_ind] = x.reshape(-1, self.out_features // len(self.enable_lora) * sum(self.enable_lora))
+        return result.view((*x.shape[:-1], self.out_features))
+
+    def train(self, mode: bool = True):
+        def T(w):
+            return w.T if self.fan_in_fan_out else w
+
+        nn.Linear.train(self, mode)
+        self.lora_A.train(mode)
+        self.lora_B.train(mode)
+        if self.merge_weights and self.merged:
+            # Make sure that the weights are not merged
+            if self.r > 0 and any(self.enable_lora):
+                delta_w = F.conv1d(
+                    self.lora_A.weight.data.unsqueeze(0),
+                    self.lora_B.weight.data.unsqueeze(-1),
+                    groups=sum(self.enable_lora),
+                ).squeeze(0)
+                self.weight.data -= self.zero_pad(T(delta_w * self.scaling))
+            self.merged = False
+
+    def eval(self):
+        def T(w):
+            return w.T if self.fan_in_fan_out else w
+
+        nn.Linear.eval(self)
+        self.lora_A.eval()
+        self.lora_B.eval()
+        if self.merge_weights and not self.merged:
+            # Merge the weights and mark it
+            if self.r > 0 and any(self.enable_lora):
+                delta_w = F.conv1d(
+                    self.lora_A.weight.data.unsqueeze(0),
+                    self.lora_B.weight.data.unsqueeze(-1),
+                    groups=sum(self.enable_lora),
+                ).squeeze(0)
+                self.weight.data += self.zero_pad(T(delta_w * self.scaling))
+            self.merged = True
+
+    def forward(self, x: torch.Tensor):
+        def T(w):
+            return w.T if self.fan_in_fan_out else w
+
+        if self.merged:
+            return F.linear(x, T(self.weight), bias=self.bias)
+        else:
+            result = F.linear(x, T(self.weight), bias=self.bias)
+            if self.r > 0:
+                after_A = self.lora_A(self.lora_dropout(x))
+                after_B = self.lora_B(after_A.transpose(-2, -1))
+                result += self.zero_pad(after_B) * self.scaling
+            return result
diff --git a/src/pet/tuners/p_tuning.py b/src/pet/tuners/p_tuning.py
index bb5a5a8..af28acf 100644
--- a/src/pet/tuners/p_tuning.py
+++ b/src/pet/tuners/p_tuning.py
@@ -14,21 +14,33 @@ class PromptEncoderReparameterizationType(str, enum.Enum):
 
 @dataclass
 class PromptEncoderConfig(PromptLearningConfig):
+    """
+    This is the configuration class to store the configuration of a :class:`~pet.PromptEncoder`.
+
+    Args:
+        encoder_reparameterization_type
+            (:obj:Union[:class:`PromptEncoderReparameterizationType`, :obj:`str`]): The type of reparameterization to
+            use.
+        encoder_hidden_size (:obj:`int`): The hidden size of the prompt encoder.
+        encoder_num_layers (:obj:`int`): The number of layers of the prompt encoder.
+        encoder_dropout (:obj:`float`): The dropout probability of the prompt encoder.
+    """
+
     encoder_reparameterization_type: Union[str, PromptEncoderReparameterizationType] = field(
         default=PromptEncoderReparameterizationType.MLP,
         metadata={"help": "How to reparameterize the prompt encoder"},
     )
     encoder_hidden_size: int = field(
         default=None,
-        metadata={"help": "The hidden size of the prompt encoder reparameterization"},
+        metadata={"help": "The hidden size of the prompt encoder"},
     )
     encoder_num_layers: int = field(
         default=2,
-        metadata={"help": "The number of layers of the prompt encoder reparameterization"},
+        metadata={"help": "The number of layers of the prompt encoder"},
     )
     encoder_dropout: float = field(
         default=0.0,
-        metadata={"help": "The dropout of the prompt encoder reparameterization"},
+        metadata={"help": "The dropout of the prompt encoder"},
     )
 
 
@@ -37,6 +49,35 @@ class PromptEncoderConfig(PromptLearningConfig):
 class PromptEncoder(torch.nn.Module):
     """
     The prompt encoder network that is used to generate the virtual token embeddings for p-tuning.
+
+    Args:
+        config (:class:`PromptEncoderConfig`): The configuration of the prompt encoder.
+
+    Example::
+
+        >>> from pet import PromptEncoder, PromptEncoderConfig >>> config = PromptEncoderConfig(
+                pet_type="P_TUNING", task_type="SEQ_2_SEQ_LM", num_virtual_tokens=20, token_dim=768,
+                num_transformer_submodules=1, num_attention_heads=12, num_layers=12,
+                encoder_reparameterization_type="MLP", encoder_hidden_size=768
+            )
+        >>> prompt_encoder = PromptEncoder(config)
+
+    Attributes:
+        embedding (:class:`~torch.nn.Embedding`): The embedding layer of the prompt encoder. mlp_head
+        (:class:`~torch.nn.Sequential`): The MLP head of the prompt encoder if `inference_mode=False`. lstm_head
+        (:class:`~torch.nn.LSTM`):
+            The LSTM head of the prompt encoder if `inference_mode=False` and `encoder_reparameterization_type="LSTM"`.
+        token_dim (:obj:`int`): The hidden embedding dimension of the base transformer model. input_size (:obj:`int`):
+        The input size of the prompt encoder. output_size (:obj:`int`): The output size of the prompt encoder.
+        hidden_size (:obj:`int`): The hidden size of the prompt encoder. total_virtual_tokens (:obj:`int`): The total
+        number of virtual tokens of the prompt encoder. encoder_type
+        (:obj:Union[:class:`PromptEncoderReparameterizationType`, :obj:`str`]):
+            The encoder type of the prompt encoder.
+
+
+    Input shape: (batch_size, total_virtual_tokens)
+
+    Output shape: (batch_size, total_virtual_tokens, token_dim)
     """
 
     def __init__(self, config):
diff --git a/src/pet/tuners/prefix_tuning.py b/src/pet/tuners/prefix_tuning.py
index d709baf..39de89d 100644
--- a/src/pet/tuners/prefix_tuning.py
+++ b/src/pet/tuners/prefix_tuning.py
@@ -8,6 +8,16 @@ from ..utils import PromptLearningConfig
 
 @dataclass
 class PrefixTuningConfig(PromptLearningConfig):
+    """
+    This is the configuration class to store the configuration of a :class:`~pet.PrefixEncoder`.
+
+    Args:
+        encoder_hidden_size (:obj: int): The hidden size of the prompt encoder.
+        prefix_projection (:obj: bool): Whether to project the prefix embeddings.
+        postprocess_past_key_value_function (:
+            obj: Optional[Callable]): The function to postprocess the past key value.
+    """
+
     encoder_hidden_size: int = field(
         default=None,
         metadata={"help": "The hidden size of the encoder"},
@@ -28,6 +38,24 @@ class PrefixEncoder(torch.nn.Module):
     r"""
     The torch.nn model to encode the prefix
 
+    Args:
+        config (:class:`PrefixTuningConfig`): The configuration of the prefix encoder.
+
+    Example::
+
+        >>> from pet import PrefixEncoder, PrefixTuningConfig >>> config = PrefixTuningConfig(
+                pet_type="PREFIX_TUNING", task_type="SEQ_2_SEQ_LM", num_virtual_tokens=20, token_dim=768,
+                num_transformer_submodules=1, num_attention_heads=12, num_layers=12, encoder_hidden_size=768
+            )
+        >>> prefix_encoder = PrefixEncoder(config)
+
+
+    Attributes:
+        embedding (:obj:`torch.nn.Embedding`): The embedding layer of the prefix encoder. trans
+        (:obj:`torch.nn.Sequential`): The two-layer MLP to transform the prefix embeddings
+            if :obj:`prefix_projection` is :obj:`True`.
+        prefix_projection (:obj:`bool`): Whether to project the prefix embeddings.
+
     Input shape: (batch_size, num_virtual_tokens)
 
     Output shape: (batch_size, num_virtual_tokens, 2*layers*hidden)
diff --git a/src/pet/tuners/prompt_tuning.py b/src/pet/tuners/prompt_tuning.py
index fdad421..5b6c313 100644
--- a/src/pet/tuners/prompt_tuning.py
+++ b/src/pet/tuners/prompt_tuning.py
@@ -15,6 +15,18 @@ class PromptTuningInit(str, enum.Enum):
 
 @dataclass
 class PromptTuningConfig(PromptLearningConfig):
+    """
+    This is the configuration class to store the configuration of a :class:`~pet.PromptEmbedding`.
+
+    Args:
+        prompt_tuning_init (:
+            obj:Union[:class:`PromptTuningInit`, :obj:`str`]): The initialization of the prompt embedding.
+        prompt_tuning_init_text (:obj: Optional[:obj:`str`]): The text to initialize the prompt embedding.
+            Only used if `prompt_tuning_init` is `TEXT`
+        tokenizer_name_or_path (:obj: Optional[:obj:`str`]): The name or path of the tokenizer.
+            Only used if `prompt_tuning_init` is `TEXT`
+    """
+
     prompt_tuning_init: Union[PromptTuningInit, str] = field(
         default=PromptTuningInit.RANDOM,
         metadata={"help": "How to initialize the prompt tuning parameters"},
@@ -34,17 +46,44 @@ class PromptTuningConfig(PromptLearningConfig):
 
 
 class PromptEmbedding(torch.nn.Module):
+    """
+    The model to encode virtual tokens into prompt embeddings.
+
+    Args:
+        config (:class:`PromptTuningConfig`): The configuration of the prompt embedding.
+        word_embeddings (:obj:`torch.nn.Module`): The word embeddings of the base transformer model.
+
+    Attributes:
+        embedding (:obj:`torch.nn.Embedding`): The embedding layer of the prompt embedding.
+
+    Example::
+
+        >>> from pet import PromptEmbedding, PromptTuningConfig >>> config = PromptTuningConfig(
+                pet_type="PROMPT_TUNING", task_type="SEQ_2_SEQ_LM", num_virtual_tokens=20, token_dim=768,
+                num_transformer_submodules=1, num_attention_heads=12, num_layers=12, prompt_tuning_init="TEXT",
+                prompt_tuning_init_text="Predict if sentiment of this review is positive, negative or neutral",
+                tokenizer_name_or_path="t5-base",
+            )
+        >>> # t5_model.shared is the word embeddings of the base model >>> prompt_embedding = PromptEmbedding(config,
+        t5_model.shared)
+
+
+    Input Shape: (batch_size, total_virtual_tokens)
+
+    Output Shape: (batch_size, total_virtual_tokens, token_dim)
+    """
+
     def __init__(self, config, word_embeddings):
         super().__init__()
 
         total_virtual_tokens = config.num_virtual_tokens * config.num_transformer_submodules
-        self.embedding = torch.nn.Embedding(total_virtual_tokens, config["token_dim"])
+        self.embedding = torch.nn.Embedding(total_virtual_tokens, config.token_dim)
         if config.prompt_tuning_init == PromptTuningInit.TEXT:
             from transformers import AutoTokenizer
 
-            self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name_or_path)
-            self.init_text = config.prompt_tuning_init_text
-            init_token_ids = self.tokenizer(self.init_text)["input_ids"]
+            tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name_or_path)
+            init_text = config.prompt_tuning_init_text
+            init_token_ids = tokenizer(init_text)["input_ids"]
             # Trim or iterate until num_text_tokens matches total_virtual_tokens
             num_text_tokens = len(init_token_ids)
             if num_text_tokens > total_virtual_tokens:
diff --git a/src/pet/utils/__init__.py b/src/pet/utils/__init__.py
index aa359db..210d711 100644
--- a/src/pet/utils/__init__.py
+++ b/src/pet/utils/__init__.py
@@ -3,5 +3,5 @@
 # module, but to preserve other warnings. So, don't check this module at all
 
 from .config import PETConfig, PETType, PromptLearningConfig, TaskType
-from .other import bloom_model_postprocess_past_key_value, shift_tokens_right
+from .other import _set_trainable, bloom_model_postprocess_past_key_value, shift_tokens_right
 from .save_and_load import get_pet_model_state_dict, set_pet_model_state_dict
diff --git a/src/pet/utils/config.py b/src/pet/utils/config.py
index d989601..c6d4588 100644
--- a/src/pet/utils/config.py
+++ b/src/pet/utils/config.py
@@ -19,7 +19,12 @@ class TaskType(str, enum.Enum):
 @dataclass
 class PETConfig:
     """
-    This is the configuration class to store the configuration of a :class:`~pet.PETModel`.
+    This is the base configuration class to store the configuration of a :class:`~pet.PETModel`.
+
+    Args:
+        pet_type (:obj:Union[:class:`~pet.utils.config.PETType`, :obj:`str`]): The type of PET method to use.
+        task_type (:obj:Union[:class:`~pet.utils.config.TaskType`, :obj:`str`]): The type of task to perform.
+        inference_mode (:obj:`bool`, defaults to :obj:`False`): Whether to use the PET model in inference mode.
     """
 
     pet_type: Union[str, PETType] = field(default=None, metadata={"help": "PET type"})
@@ -29,8 +34,22 @@ class PETConfig:
 
 @dataclass
 class PromptLearningConfig(PETConfig):
+    """
+    This is the base configuration class to store the configuration of a :obj:Union[:class:`~pet.PrefixTuning`,
+    :class:`~pet.PromptEncoder`, :class:`~pet.PromptTuning`].
+
+    Args:
+        num_virtual_tokens (:obj:`int`): The number of virtual tokens to use.
+        token_dim (:obj:`int`): The hidden embedding dimension of the base transformer model.
+        num_transformer_submodules (:obj:`int`): The number of transformer submodules in the base transformer model.
+        num_attention_heads (:obj:`int`): The number of attention heads in the base transformer model.
+        num_layers (:obj:`int`): The number of layers in the base transformer model.
+    """
+
     num_virtual_tokens: int = field(default=None, metadata={"help": "Number of virtual tokens"})
-    token_dim: int = field(default=None, metadata={"help": "Dimension of virtual tokens"})
+    token_dim: int = field(
+        default=None, metadata={"help": "The hidden embedding dimension of the base transformer model"}
+    )
     num_transformer_submodules: Optional[int] = field(default=1, metadata={"help": "Number of transformer submodules"})
     num_attention_heads: Optional[int] = field(default=None, metadata={"help": "Number of attention heads"})
     num_layers: Optional[int] = field(default=None, metadata={"help": "Number of transformer layers"})
diff --git a/src/pet/utils/other.py b/src/pet/utils/other.py
index ac8b329..b7a81d5 100644
--- a/src/pet/utils/other.py
+++ b/src/pet/utils/other.py
@@ -19,6 +19,11 @@ def bloom_model_postprocess_past_key_value(past_key_values):
 def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
     """
     Shift input ids one token to the right.
+
+    Args:
+        input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`): input ids
+        pad_token_id (:obj:`int`): The id of the `padding` token.
+        decoder_start_token_id (:obj:`int`): The id of the `start` token.
     """
     shifted_input_ids = input_ids.new_zeros(input_ids.shape)
     shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
@@ -30,3 +35,47 @@ def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start
     shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
 
     return shifted_input_ids
+
+
+def _set_trainable(model):
+    if model.modules_to_save is not None:
+        for name, param in model.named_parameters():
+            if any(module_name in name for module_name in model.modules_to_save):
+                param.requires_grad = True
+            else:
+                param.requires_grad = False
+
+
+def fsdp_auto_wrap_policy(model):
+    import functools
+    import os
+
+    from accelerate import FullyShardedDataParallelPlugin
+    from torch.distributed.fsdp.wrap import _or_policy, lambda_auto_wrap_policy, transformer_auto_wrap_policy
+
+    from ..tuners import PrefixEncoder, PromptEmbedding, PromptEncoder
+
+    def lambda_policy_fn(module):
+        if (
+            len(list(module.named_children())) == 0
+            and getattr(module, "weight", None) is not None
+            and module.weight.requires_grad
+        ):
+            return True
+        return False
+
+    lambda_policy = functools.partial(lambda_auto_wrap_policy, lambda_fn=lambda_policy_fn)
+    transformer_wrap_policy = functools.partial(
+        transformer_auto_wrap_policy,
+        transformer_layer_cls=(
+            PrefixEncoder,
+            PromptEncoder,
+            PromptEmbedding,
+            FullyShardedDataParallelPlugin.get_module_class_from_name(
+                model, os.environ.get("FSDP_TRANSFORMER_CLS_TO_WRAP", "")
+            ),
+        ),
+    )
+
+    auto_wrap_policy = functools.partial(_or_policy, policies=[lambda_policy, transformer_wrap_policy])
+    return auto_wrap_policy
diff --git a/src/pet/utils/save_and_load.py b/src/pet/utils/save_and_load.py
index e2eea05..1418dfa 100644
--- a/src/pet/utils/save_and_load.py
+++ b/src/pet/utils/save_and_load.py
@@ -4,6 +4,13 @@ from .config import PETType
 
 
 def get_pet_model_state_dict(model):
+    """
+    Get the state dict of the PET model.
+
+    Args:
+        model (:obj:`PETModel`): The PET model.
+    """
+
     if model.pet_config.pet_type == PETType.LORA:
         return lora_state_dict(model)
     else:
@@ -19,6 +26,14 @@ def get_pet_model_state_dict(model):
 
 
 def set_pet_model_state_dict(model, pet_model_state_dict):
+    """
+    Set the state dict of the PET model.
+
+    Args:
+        model (:obj:`PETModel`): The PET model.
+        pet_model_state_dict (:obj:`dict`): The state dict of the PET model.
+    """
+
     model.load_state_dict(pet_model_state_dict, strict=False)
     if model.pet_config.pet_type != PETType.LORA:
         model.prompt_encoder.embedding.load_state_dict(