diff --git a/Dockerfile b/Dockerfile
index d933bff..5d671c2 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -15,6 +15,8 @@ FROM python:3.12-slim AS backend-kubeflow-wheel
 
 WORKDIR /src
 
+RUN pip install setuptools wheel
+
 ARG BACKEND_LIB=/kf/kubeflow/components/crud-web-apps/common/backend
 COPY --from=fetch-kubeflow-kubeflow $BACKEND_LIB .
 RUN python setup.py sdist bdist_wheel
diff --git a/frontend/src/app/shared/utils.ts b/frontend/src/app/shared/utils.ts
index f637e92..6b78ce8 100644
--- a/frontend/src/app/shared/utils.ts
+++ b/frontend/src/app/shared/utils.ts
@@ -143,14 +143,7 @@ export function getPredictorExtensionSpec(
   predictor: PredictorSpec,
 ): PredictorExtensionSpec {
   if (predictor.model) {
-    if (
-      Object.values(PredictorType).includes(
-        predictor.model?.modelFormat.name as PredictorType,
-      )
-    ) {
-      const spec = predictor.model;
-      return spec;
-    }
+    return predictor.model;
   } else {
     for (const predictorType of Object.values(PredictorType)) {
       if (predictorType in predictor) {