Use the defined mode for input/output types

replicate · Oct 28, 2024 · 5cf3ace · 5cf3ace
1 parent 2840fc7
commit 5cf3ace
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/cog/server/http.py b/python/cog/server/http.py
@@ -151,7 +151,7 @@ async def start_shutdown() -> Any:
         return JSONResponse({}, status_code=200)
 
     try:
-        InputType, OutputType = cog_config.get_predictor_types(mode=Mode.PREDICT)
+        InputType, OutputType = cog_config.get_predictor_types(mode=mode)
     except Exception:  # pylint: disable=broad-exception-caught
         msg = "Error while loading predictor:\n\n" + traceback.format_exc()
         add_setup_failed_routes(app, started_at, msg)