Update app.py

app.py

@@ -1,4 +1,4 @@
-
+import spaces
 import streamlit as st
 from transformers import pipeline
 from PIL import Image
@@ -7,6 +7,7 @@ from peft import AutoPeftModelForCausalLM
 from transformers import AutoTokenizer
 import torch
 
+@spaces.GPU
 def main():
 
     def get_tokens_as_list(word_list):
@@ -59,8 +60,9 @@ def main():
     
     
     model_id = "r1208/c4ai-command-r-v01-4bit_32r"
+
     
-    model = AutoPeftModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, use_auth_token=hf_token, device_map = "auto")
+    model = AutoPeftModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, use_auth_token=hf_token)
     tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=hf_token)
     
     tokenizer_with_prefix_space = AutoTokenizer.from_pretrained(model_id, add_prefix_space=True, use_auth_token=hf_token)