diff --git a/README.md b/README.md
index c8eb43b..df03330 100644
--- a/README.md
+++ b/README.md
@@ -45,6 +45,7 @@ Hardware: Single A100 80GB GPU with CPU RAM above 64GB
 | --------- | ---- | ---- | ---- |
 | bigscience/T0_3B (3B params) | 47.14GB GPU / 2.96GB CPU  | 14.4GB GPU / 2.96GB CPU | 9.8GB GPU / 17.8GB CPU |
 | bigscience/mt0-xxl (12B params) | OOM GPU | 56GB GPU / 3GB CPU | 22GB GPU / 52GB CPU |
+| bigscience/bloomz-7b1 (7B params) | OOM GPU | 32GB GPU / 3.8GB CPU | 18.1GB GPU / 35GB CPU |
 
 Performance of PET-LoRA tuned `bigscience/T0_3B` on `ought/raft/twitter_complaints` leaderboard. 
 A point to note is that we didn't try to sequeeze performance by playing around with input instruction templates, LoRA hyperparams and other training related hyperparams. Also, we didn't use the larger 13B mt0-xxl model.
diff --git a/examples/causal_language_modeling/pet_lora_clm_accelerate_ds_zero3_offload.py b/examples/causal_language_modeling/pet_lora_clm_accelerate_ds_zero3_offload.py
new file mode 100644
index 0000000..d919553
--- /dev/null
+++ b/examples/causal_language_modeling/pet_lora_clm_accelerate_ds_zero3_offload.py
@@ -0,0 +1,351 @@
+import gc
+import os
+import sys
+import threading
+
+import numpy as np
+import torch
+from accelerate import Accelerator
+from torch.utils.data import DataLoader
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    default_data_collator,
+    get_linear_schedule_with_warmup,
+    set_seed,
+)
+
+import psutil
+from datasets import load_dataset
+from pet import LoRAConfig, TaskType, get_pet_model, get_pet_model_state_dict
+from tqdm import tqdm
+
+
+def levenshtein_distance(str1, str2):
+    # TC: O(N^2)
+    # SC: O(N^2)
+    if str1 == str2:
+        return 0
+    num_rows = len(str1) + 1
+    num_cols = len(str2) + 1
+    dp_matrix = np.empty((num_rows, num_cols))
+    dp_matrix[0, :] = range(num_cols)
+    dp_matrix[:, 0] = range(num_rows)
+
+    for i in range(1, num_rows):
+        for j in range(1, num_cols):
+            if str1[i - 1] == str2[j - 1]:
+                dp_matrix[i, j] = dp_matrix[i - 1, j - 1]
+            else:
+                dp_matrix[i, j] = min(dp_matrix[i - 1, j - 1], dp_matrix[i - 1, j], dp_matrix[i, j - 1]) + 1
+
+    return dp_matrix[num_rows - 1, num_cols - 1]
+
+
+def get_closest_label(eval_pred, classes):
+    min_id = sys.maxsize
+    min_edit_distance = sys.maxsize
+    for i, class_label in enumerate(classes):
+        edit_distance = levenshtein_distance(eval_pred.strip(), class_label)
+        if edit_distance < min_edit_distance:
+            min_id = i
+            min_edit_distance = edit_distance
+    return classes[min_id]
+
+
+# Converting Bytes to Megabytes
+def b2mb(x):
+    return int(x / 2**20)
+
+
+# This context manager is used to track the peak memory usage of the process
+class TorchTracemalloc:
+    def __enter__(self):
+        gc.collect()
+        torch.cuda.empty_cache()
+        torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
+        self.begin = torch.cuda.memory_allocated()
+        self.process = psutil.Process()
+
+        self.cpu_begin = self.cpu_mem_used()
+        self.peak_monitoring = True
+        peak_monitor_thread = threading.Thread(target=self.peak_monitor_func)
+        peak_monitor_thread.daemon = True
+        peak_monitor_thread.start()
+        return self
+
+    def cpu_mem_used(self):
+        """get resident set size memory for the current process"""
+        return self.process.memory_info().rss
+
+    def peak_monitor_func(self):
+        self.cpu_peak = -1
+
+        while True:
+            self.cpu_peak = max(self.cpu_mem_used(), self.cpu_peak)
+
+            # can't sleep or will not catch the peak right (this comment is here on purpose)
+            # time.sleep(0.001) # 1msec
+
+            if not self.peak_monitoring:
+                break
+
+    def __exit__(self, *exc):
+        self.peak_monitoring = False
+
+        gc.collect()
+        torch.cuda.empty_cache()
+        self.end = torch.cuda.memory_allocated()
+        self.peak = torch.cuda.max_memory_allocated()
+        self.used = b2mb(self.end - self.begin)
+        self.peaked = b2mb(self.peak - self.begin)
+
+        self.cpu_end = self.cpu_mem_used()
+        self.cpu_used = b2mb(self.cpu_end - self.cpu_begin)
+        self.cpu_peaked = b2mb(self.cpu_peak - self.cpu_begin)
+        # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")
+
+
+def main():
+    accelerator = Accelerator()
+    model_name_or_path = "bigscience/bloomz-7b1"
+    dataset_name = "twitter_complaints"
+    pet_config = LoRAConfig(task_type=TaskType.CAUSAL_LM, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1)
+    checkpoint_name = (
+        f"{dataset_name}_{model_name_or_path}_{pet_config.pet_type}_{pet_config.task_type}_v1.pt".replace("/", "_")
+    )
+    text_column = "Tweet text"
+    label_column = "text_label"
+    lr = 3e-3
+    num_epochs = 20
+    batch_size = 8
+    seed = 42
+    max_length = 64
+    set_seed(seed)
+
+    dataset = load_dataset("ought/raft", dataset_name)
+    classes = [k.replace("_", " ") for k in 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):
+        batch_size = len(examples[text_column])
+        inputs = [f"{text_column} : {x} Label : " for x in examples[text_column]]
+        targets = [str(x) for x in examples[label_column]]
+        model_inputs = tokenizer(inputs)
+        labels = tokenizer(targets)
+        for i in range(batch_size):
+            sample_input_ids = model_inputs["input_ids"][i]
+            label_input_ids = labels["input_ids"][i] + [tokenizer.pad_token_id]
+            model_inputs["input_ids"][i] = sample_input_ids + label_input_ids
+            labels["input_ids"][i] = [-100] * len(sample_input_ids) + label_input_ids
+            model_inputs["attention_mask"][i] = [1] * len(model_inputs["input_ids"][i])
+        for i in range(batch_size):
+            sample_input_ids = model_inputs["input_ids"][i]
+            label_input_ids = labels["input_ids"][i]
+            model_inputs["input_ids"][i] = [tokenizer.pad_token_id] * (
+                max_length - len(sample_input_ids)
+            ) + sample_input_ids
+            model_inputs["attention_mask"][i] = [0] * (max_length - len(sample_input_ids)) + model_inputs[
+                "attention_mask"
+            ][i]
+            labels["input_ids"][i] = [-100] * (max_length - len(sample_input_ids)) + label_input_ids
+            model_inputs["input_ids"][i] = torch.tensor(model_inputs["input_ids"][i][:max_length])
+            model_inputs["attention_mask"][i] = torch.tensor(model_inputs["attention_mask"][i][:max_length])
+            labels["input_ids"][i] = torch.tensor(labels["input_ids"][i][:max_length])
+        model_inputs["labels"] = labels["input_ids"]
+        return model_inputs
+
+    def test_preprocess_function(examples):
+        batch_size = len(examples[text_column])
+        inputs = [f"{text_column} : {x} Label : " for x in examples[text_column]]
+        model_inputs = tokenizer(inputs)
+        # print(model_inputs)
+        for i in range(batch_size):
+            sample_input_ids = model_inputs["input_ids"][i]
+            model_inputs["input_ids"][i] = [tokenizer.pad_token_id] * (
+                max_length - len(sample_input_ids)
+            ) + sample_input_ids
+            model_inputs["attention_mask"][i] = [0] * (max_length - len(sample_input_ids)) + model_inputs[
+                "attention_mask"
+            ][i]
+            model_inputs["input_ids"][i] = torch.tensor(model_inputs["input_ids"][i][:max_length])
+            model_inputs["attention_mask"][i] = torch.tensor(model_inputs["attention_mask"][i][:max_length])
+        return model_inputs
+
+    with accelerator.main_process_first():
+        processed_datasets = dataset.map(
+            preprocess_function,
+            batched=True,
+            num_proc=1,
+            remove_columns=dataset["train"].column_names,
+            load_from_cache_file=True,
+            desc="Running tokenizer on dataset",
+        )
+    accelerator.wait_for_everyone()
+
+    train_dataset = processed_datasets["train"]
+
+    with accelerator.main_process_first():
+        processed_datasets = dataset.map(
+            test_preprocess_function,
+            batched=True,
+            num_proc=1,
+            remove_columns=dataset["train"].column_names,
+            load_from_cache_file=False,
+            desc="Running tokenizer on dataset",
+        )
+    eval_dataset = processed_datasets["train"]
+    test_dataset = processed_datasets["test"]
+
+    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
+    )
+    test_dataloader = DataLoader(
+        test_dataset, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True
+    )
+
+    print(next(iter(train_dataloader)))
+
+    # creating model
+    model = AutoModelForCausalLM.from_pretrained(model_name_or_path)
+    model = get_pet_model(model, pet_config)
+    model.print_trainable_parameters()
+
+    # optimizer
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
+
+    # lr scheduler
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer=optimizer,
+        num_warmup_steps=0,
+        num_training_steps=(len(train_dataloader) * num_epochs),
+    )
+
+    model, train_dataloader, eval_dataloader, test_dataloader, optimizer, lr_scheduler = accelerator.prepare(
+        model, train_dataloader, eval_dataloader, test_dataloader, optimizer, lr_scheduler
+    )
+    accelerator.print(model)
+
+    is_ds_zero_3 = False
+    if getattr(accelerator.state, "deepspeed_plugin", None):
+        is_ds_zero_3 = accelerator.state.deepspeed_plugin.zero_stage == 3
+
+    for epoch in range(num_epochs):
+        with TorchTracemalloc() as tracemalloc:
+            model.train()
+            total_loss = 0
+            for step, batch in enumerate(tqdm(train_dataloader)):
+                outputs = model(**batch)
+                loss = outputs.loss
+                total_loss += loss.detach().float()
+                accelerator.backward(loss)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
+        accelerator.print("GPU Memory before entering the train : {}".format(b2mb(tracemalloc.begin)))
+        accelerator.print("GPU Memory consumed at the end of the train (end-begin): {}".format(tracemalloc.used))
+        accelerator.print("GPU Peak Memory consumed during the train (max-begin): {}".format(tracemalloc.peaked))
+        accelerator.print(
+            "GPU Total Peak Memory consumed during the train (max): {}".format(
+                tracemalloc.peaked + b2mb(tracemalloc.begin)
+            )
+        )
+
+        accelerator.print("CPU Memory before entering the train : {}".format(b2mb(tracemalloc.cpu_begin)))
+        accelerator.print("CPU Memory consumed at the end of the train (end-begin): {}".format(tracemalloc.cpu_used))
+        accelerator.print("CPU Peak Memory consumed during the train (max-begin): {}".format(tracemalloc.cpu_peaked))
+        accelerator.print(
+            "CPU Total Peak Memory consumed during the train (max): {}".format(
+                tracemalloc.cpu_peaked + b2mb(tracemalloc.cpu_begin)
+            )
+        )
+        train_epoch_loss = total_loss / len(eval_dataloader)
+        train_ppl = torch.exp(train_epoch_loss)
+        accelerator.print(f"{epoch=}: {train_ppl=} {train_epoch_loss=}")
+
+        model.eval()
+        eval_preds = []
+        with TorchTracemalloc() as tracemalloc:
+            for _, batch in enumerate(tqdm(eval_dataloader)):
+                batch = {k: v for k, v in batch.items() if k != "labels"}
+                with torch.no_grad():
+                    outputs = accelerator.unwrap_model(model).generate(
+                        **batch, synced_gpus=is_ds_zero_3, max_new_tokens=10
+                    )  # synced_gpus=True for DS-stage 3
+                preds = outputs[:, max_length:].detach().cpu().numpy()
+                eval_preds.extend(tokenizer.batch_decode(preds, skip_special_tokens=True))
+
+        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
+        accelerator.print("GPU Memory before entering the eval : {}".format(b2mb(tracemalloc.begin)))
+        accelerator.print("GPU Memory consumed at the end of the eval (end-begin): {}".format(tracemalloc.used))
+        accelerator.print("GPU Peak Memory consumed during the eval (max-begin): {}".format(tracemalloc.peaked))
+        accelerator.print(
+            "GPU Total Peak Memory consumed during the eval (max): {}".format(
+                tracemalloc.peaked + b2mb(tracemalloc.begin)
+            )
+        )
+
+        accelerator.print("CPU Memory before entering the eval : {}".format(b2mb(tracemalloc.cpu_begin)))
+        accelerator.print("CPU Memory consumed at the end of the eval (end-begin): {}".format(tracemalloc.cpu_used))
+        accelerator.print("CPU Peak Memory consumed during the eval (max-begin): {}".format(tracemalloc.cpu_peaked))
+        accelerator.print(
+            "CPU Total Peak Memory consumed during the eval (max): {}".format(
+                tracemalloc.cpu_peaked + b2mb(tracemalloc.cpu_begin)
+            )
+        )
+
+        correct = 0
+        total = 0
+        for pred, true in zip(eval_preds, dataset["train"][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.print(f"{dataset['train'][label_column][:10]=}")
+
+    model.eval()
+    test_preds = []
+    for _, batch in enumerate(tqdm(test_dataloader)):
+        batch = {k: v for k, v in batch.items() if k != "labels"}
+        with torch.no_grad():
+            outputs = accelerator.unwrap_model(model).generate(
+                **batch, synced_gpus=is_ds_zero_3, max_new_tokens=10
+            )  # synced_gpus=True for DS-stage 3
+        test_preds.extend(
+            tokenizer.batch_decode(outputs[:, max_length:].detach().cpu().numpy(), skip_special_tokens=True)
+        )
+
+    test_preds_cleaned = []
+    for _, pred in enumerate(test_preds):
+        test_preds_cleaned.append(get_closest_label(pred, classes))
+
+    test_df = dataset["test"].to_pandas()
+    test_df[label_column] = test_preds_cleaned
+    test_df["text_labels_orig"] = test_preds
+    accelerator.print(test_df[[text_column, label_column]].sample(20))
+
+    pred_df = test_df[["ID", label_column]]
+    pred_df.columns = ["ID", "Label"]
+
+    os.makedirs(f"data/{dataset_name}", exist_ok=True)
+    pred_df.to_csv(f"data/{dataset_name}/predictions.csv", index=False)
+
+    accelerator.wait_for_everyone()
+    accelerator.save(get_pet_model_state_dict(model, state_dict=accelerator.get_state_dict(model)), checkpoint_name)
+    accelerator.wait_for_everyone()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/conditional_generation/pet_lora_seq2seq_accelerate_ds_zero3_offload.py b/examples/conditional_generation/pet_lora_seq2seq_accelerate_ds_zero3_offload.py
index b304ce7..e3f3b31 100644
--- a/examples/conditional_generation/pet_lora_seq2seq_accelerate_ds_zero3_offload.py
+++ b/examples/conditional_generation/pet_lora_seq2seq_accelerate_ds_zero3_offload.py
@@ -107,7 +107,9 @@ def main():
     pet_config = LoRAConfig(
         task_type=TaskType.SEQ_2_SEQ_LM, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1
     )
-    checkpoint_name = f"{dataset_name}_{pet_config.pet_type}_{pet_config.task_type}_v1.pt".replace("/", "_")
+    checkpoint_name = (
+        f"{dataset_name}_{model_name_or_path}_{pet_config.pet_type}_{pet_config.task_type}_v1.pt".replace("/", "_")
+    )
     text_column = "Tweet text"
     label_column = "text_label"
     lr = 3e-3
diff --git a/src/pet/__init__.py b/src/pet/__init__.py
index 31a3df7..2a48832 100644
--- a/src/pet/__init__.py
+++ b/src/pet/__init__.py
@@ -33,4 +33,5 @@ from .utils import (
     get_pet_model_state_dict,
     set_pet_model_state_dict,
     shift_tokens_right,
+    pet_model_load_and_dispatch,
 )
diff --git a/src/pet/utils/__init__.py b/src/pet/utils/__init__.py
index 210d711..af3188e 100644
--- a/src/pet/utils/__init__.py
+++ b/src/pet/utils/__init__.py
@@ -4,4 +4,4 @@
 
 from .config import PETConfig, PETType, PromptLearningConfig, TaskType
 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
+from .save_and_load import get_pet_model_state_dict, set_pet_model_state_dict, pet_model_load_and_dispatch