Update main.py

main.py

@@ -2,8 +2,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from flask import Flask, request, jsonify
 
 device = "cuda" # the device to load the model onto
-from transformers import set_cache_dir
-set_cache_dir("/code/.cache/huggingface")
+
 
 model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
 tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")