Update configs.py

configs.py

@@ -17,8 +17,9 @@ model_info = {
     'LLAMA2-7B': dict(model_path='meta-llama/Llama-2-7b-chat-hf', token=os.environ['hf_token'], 
                       original_prompt_template='<s>{prompt}',
                       interpretation_prompt_template='<s>[INST] [X] [/INST] {prompt}',
-                      layers_format=llama_layers_format), # , load_in_8bit=True
-    'LLAMA2-13B': dict(model_path='meta-llama/Llama-2-13b-chat-hf', token=os.environ['hf_token'], 
+                      layers_format=llama_layers_format),
+    'LLAMA2-13B': dict(model_path='meta-llama/Llama-2-13b-chat-hf', 
+                      token=os.environ['hf_token'], device_map='cuda', load_in_8bit=True, 
                       original_prompt_template='<s>{prompt}',
                       interpretation_prompt_template='<s>[INST] [X] [/INST] {prompt}',
                       layers_format=llama_layers_format),