added disk_offload function

app.py

@@ -26,6 +26,7 @@ locale.getpreferredencoding = getpreferredencoding
 
 from torch import cuda, bfloat16
 import transformers
+from accelerate import disk_offload
 
 # Model used
 model_id = 'meta-llama/Llama-2-7b-chat-hf'
@@ -52,14 +53,14 @@ model_config = transformers.AutoConfig.from_pretrained(
 )
 
 # Downloading and Initializing the model
-model = transformers.AutoModelForCausalLM.from_pretrained(
+model = disk_offload(transformers.AutoModelForCausalLM.from_pretrained(
     model_id,
     trust_remote_code=True,
     config=model_config,
     # quantization_config=bnb_config,
     device_map='auto',
     token=hf_auth
-)
+))
 
 # enable evaluation mode to allow model inference
 model.eval()