trying out the sharded vesion

model.py

@@ -7,7 +7,8 @@ from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIt
 from huggingface_hub import login
 
 # model_id = 'meta-llama/Llama-2-13b-chat-hf'
-model_id = 'meta-llama/Llama-2-7b-chat-hf'
+#model_id = 'meta-llama/Llama-2-7b-chat-hf'
+model_id = 'Trelis/Llama-2-7b-chat-hf-sharded-bf16'
 
 if torch.cuda.is_available():
     tok = os.environ['HF_TOKEN']