rollback

data/models/llama3-1-8b.py

@@ -16,12 +16,12 @@ pipeline = transformers.pipeline(
     "text-generation",
     model=model_id,
     model_kwargs={"torch_dtype": torch.bfloat16},
-    device="cpu",
+    device="gpu",
     token=api_key
 )
 
 @app.route('/chat', methods=['POST'])
-#@spaces.GPU(enable_queue=True)
+@spaces.GPU(enable_queue=True)
 def chat_completion():
     data = request.json