ZERO

app.py

@@ -7,6 +7,15 @@ import time
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
+zero = torch.Tensor([0]).cuda()
+print(zero.device) # <-- 'cpu' 🤔
+
+
+@spaces.GPU
+def greet(n):
+    print(zero.device) # <-- 'cuda:0' 🤗
+    return f"Hello {zero + n} Tensor"
+
 def capture_logs(log_body, log_file, uuid_label):
     logger = logging.getLogger('MyApp')
     logger.setLevel(logging.INFO)
@@ -60,6 +69,7 @@ else:
 
 
 # Function to handle user input and generate a response
+@spaces.GPU
 def chatbot_response(query, tokens, top_k, top_p):
     uuid_label = str(uuid.uuid4())
 

requirements.txt

@@ -6,4 +6,5 @@ minijinja
 torch --extra-index-url https://download.pytorch.org/whl/cu118
 torchvision
 torchaudio
-accelerate==0.26.0
+accelerate==0.26.0
+spaces