Update app.py

app.py

@@ -116,15 +116,13 @@ snapshot_download(
   allow_patterns=["*.json", "*.safetensors", "*.bin", "*.py", "*.md", "*.txt"],
 )
 
-@spaces.GPU(duration=120)
-def load_model(save_dir):
-    model = LLM(model=save_dir,
+device = "cuda"
+model = LLM(model=save_dir,
                 max_num_seqs=512,
                 # tensor_parallel_size=8,
                 enable_prefix_caching=True, 
-                gpu_memory_utilization=0.95)
+                gpu_memory_utilization=0.95).to(device)
 
-model = load_model(save_dir)
 decoding_params = BeamSearchParams(beam_width=4,
                                    max_tokens=512)