Add device check

fpem123 · Apr 28, 2021 · e4ba7a0 · e4ba7a0
1 parent aceee65
commit e4ba7a0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/openchat/agents/blender.py b/openchat/agents/blender.py
@@ -55,6 +55,6 @@ def set_options(self, name, device):
         datapath = option.get("datapath")
         option['model_file'] = modelzoo_path(datapath, name)
         option["override"] = {
-            "no_cuda": False,
+            "no_cuda": False if "cuda" in device else True,
         }
         return option
diff --git a/openchat/envs/various_web_demo_env.py b/openchat/envs/various_web_demo_env.py
@@ -80,6 +80,7 @@ def generate(user_id, bot_id, user_message, topic, agent: str):
                 # get agent obj
                 agent = agent.upper()
                 agent_obj = self.agents[agent]
+                print(agent_obj.model.device)
 
                 # max hold 50 user for memory
                 if len(self.users) > self.max_hold_user:
@@ -213,7 +214,8 @@ def send(user_id):
                 traceback.print_exc()
                 return {'output': 'Sorry, there was an error.'}, 500
 
-        self.app.run(host='0.0.0.0', port=80)
+        from waitress import serve
+        serve(self.app, host='0.0.0.0', port=80)
 
     def pre_dialog_for_special_tasks(self, agent, user_id, bot_id, topic):
         if isinstance(agent, ConvAI2Agent):