def download_model_to_folder():
    from huggingface_hub import snapshot_download

    snapshot_download(
-        "meta-llama/Llama-2-13b-chat-hf",
+        "meta-llama/Llama-2-70b-chat-hf",
        local_dir="/model",
        token=os.environ["HUGGINGFACE_TOKEN"],
    )

image = (
    Image.from_dockerhub("nvcr.io/nvidia/pytorch:22.12-py3")
    .pip_install("torch==2.0.1", index_url="https://download.pytorch.org/whl/cu118")
+    # Pin vLLM to 8/2/2023
+    .pip_install("vllm @ git+https://github.com/vllm-project/vllm.git@79af7e96a0e2fc9f340d1939192122c3ae38ff17")
-    # Pin vLLM to 07/19/2023
-    .pip_install("vllm @ git+https://github.com/vllm-project/vllm.git@bda41c70ddb124134935a90a0d51304d2ac035e8")
    # Use the barebones hf-transfer package for maximum download speeds. No progress bar, but expect 700MB/s.
-    .pip_install("hf-transfer~=0.1")
+     #Force a rebuild to invalidate the cache (you can remove `force_build=True` after the first time)
+    .pip_install("hf-transfer~=0.1", force_build=True)
    .run_function(
        download_model_to_folder,
        secret=Secret.from_name("huggingface"),
        timeout=60 * 20)
)
...

-@stub.cls(gpu="A100", secret=Secret.from_name("huggingface"))
+# You need a minimum of 4 A100s that are the 40GB version
+@stub.cls(gpu=gpu.A100(count=4, memory=40), secret=Secret.from_name("huggingface"))
class Model:
    def __enter__(self):
        from vllm import LLM

        # Load the model. Tip: MPT models may require `trust_remote_code=true`.
-       self.llm = LLM(MODEL_DIR)
+       self.llm = LLM(MODEL_DIR, tensor_parallel_size=4)
...