Update app.py

app.py

@@ -1,5 +1,11 @@
+import torch
 import gradio as gr
 
+from gen import get_pretrained_models, get_output, setup_model_parallel
+
+torch.cuda.set_device(0)
+torch.cuda.empty_cache()
+
 def chat(user_input):
     bot_response = "hello world"
     response = ""