loading the model in CPU mode

model.py

@@ -3,29 +3,37 @@ from threading import Thread
 from typing import Iterator
 
 import torch
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
 from huggingface_hub import login
 
 model_id = 'jaymojnidar/Llama-2-7b-chat-hf-sharded-bf16-5GBMAX'
 
 if not torch.cuda.is_available():
     tok = os.environ['HF_TOKEN']
+    device_map = {
+    "transformer.word_embeddings": 0,
+    "transformer.word_embeddings_layernorm": 0,
+    "lm_head": "cpu",
+    "transformer.h": 0,
+    "transformer.ln_f": 0,
+}
     login(new_session=True,
           write_permission=False,
           token=tok
           
           #, token="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"
           )
-    
+    quantization_config = BitsAndBytesConfig(llm_int8_enable_fp32_cpu_offload=True)
+
     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,
-        load_in_4bit=True,
-        device_map='auto',
+        device_map=device_map,
         use_auth_token=True
     )
 else: