loading the model back in GPU mode

app.py

@@ -30,10 +30,10 @@ As a derivate work of [Llama-2-13b-chat](https://huggingface.co/meta-llama/Llama
 this demo is governed by the original [license](https://huggingface.co/spaces/huggingface-projects/llama-2-13b-chat/blob/main/LICENSE.txt) and [acceptable use policy](https://huggingface.co/spaces/huggingface-projects/llama-2-13b-chat/blob/main/USE_POLICY.md).
 """
 
-'''
+
 if not torch.cuda.is_available():
     DESCRIPTION += '\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>'
-'''
+
 
 def clear_and_save_textbox(message: str) -> tuple[str, str]:
     return '', message

model.py

@@ -3,41 +3,29 @@ from threading import Thread
 from typing import Iterator
 
 import torch
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from huggingface_hub import login
 
 model_id = 'jaymojnidar/Llama-2-7b-chat-hf-sharded-bf16-5GBMAX'
 
-if not torch.cuda.is_available():
+if torch.cuda.is_available():
     tok = os.environ['HF_TOKEN']
-    
-    device_map = {
-        "transformer.word_embeddings": "cpu",
-        "transformer.word_embeddings_layernorm": "cpu",
-        "lm_head": "cpu",
-        "transformer.h": "cpu",
-        "transformer.ln_f": "cpu",
-        "model.layers": "cpu",
-        "model.norm": "cpu",
-        }
-  
     login(new_session=True,
           write_permission=False,
           token=tok
           
           #, token="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"
           )
-    quantization_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload=True, load_in_8bit=True,llm_int8_threshold=200.0)
-
+    
     config = AutoConfig.from_pretrained(model_id,
         use_auth_token=True)
     config.pretraining_tp = 1
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         config=config,
-        quantization_config=quantization_config,
         torch_dtype=torch.float16,
-        device_map=device_map,
+        load_in_8bit=True,
+        device_map='auto',
         use_auth_token=True
     )
 else:
@@ -73,7 +61,7 @@ def run(message: str,
         top_p: float = 0.95,
         top_k: int = 50) -> Iterator[str]:
     prompt = get_prompt(message, chat_history, system_prompt)
-    inputs = tokenizer([prompt], return_tensors='pt', add_special_tokens=False).to(torch.device) #.to('cuda')
+    inputs = tokenizer([prompt], return_tensors='pt', add_special_tokens=False).to('cuda')
 
     streamer = TextIteratorStreamer(tokenizer,
                                     timeout=10.,