Update README.md

README.md

@@ -62,11 +62,11 @@ instruction = """"Help me set up my daily to-do list!""""
 ```
 ```python
 from peft import PeftModel, PeftConfig
-from transformers import AutoModelForCausalLM
+from transformers import AutoModelForCausalLM,AutoTokenizer
 
-config = PeftConfig.from_pretrained("SwastikM/Llama-2-7B-Chat-text2code")      #PEFT Config
-model = AutoModelForCausalLM.from_pretrained("TheBloke/Llama-2-7b-Chat-GPTQ")  #Loading the Base Model
-model = PeftModel.from_pretrained(model, "SwastikM/Llama-2-7B-Chat-text2code") #Combining Trained Adapter with Base Model
+config = PeftConfig.from_pretrained("SwastikM/Llama-2-7B-Chat-text2code")                        #PEFT Config
+model = AutoModelForCausalLM.from_pretrained("TheBloke/Llama-2-7b-Chat-GPTQ",device_map='auto')  #Loading the Base Model
+model = PeftModel.from_pretrained(model, "SwastikM/Llama-2-7B-Chat-text2code")                   #Combining Trained Adapter with Base Model
 tokenizer = AutoTokenizer.from_pretrained("SwastikM/Llama-2-7B-Chat-text2code")
 
 inputs = tokenizer(instruction, return_tensors="pt").input_ids.to('cuda')