Update app.py

app.py

@@ -1,12 +1,20 @@
 from peft import PeftModel
 from transformers import LLaMATokenizer, LLaMAForCausalLM, GenerationConfig
+import torch
+n_gpus = torch.cuda.device_count()
+max_memory = {i: max_memory for i in range(n_gpus)}
+
+print(f'Max memory : {max_memory}')
 
 tokenizer = LLaMATokenizer.from_pretrained("decapoda-research/llama-7b-hf")
+max_memory = '40GB'
+
 model = LLaMAForCausalLM.from_pretrained(
     "decapoda-research/llama-7b-hf",
     load_in_8bit=True,
-    device_map="auto",
+    device_map="auto",max_memory=max_memory
 )
+
 model = PeftModel.from_pretrained(model, "tloen/alpaca-lora-7b")
 
 def generate_prompt(instruction, input=None):