diff --git a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
index 8614a0685..3d36b4e26 100644
--- a/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
+++ b/src/frontend/src/CustomNodes/GenericNode/components/parameterComponent/index.tsx
@@ -16,6 +16,7 @@ import { PopUpContext } from "../../../../contexts/popUpContext";
 import { TabsContext } from "../../../../contexts/tabsContext";
 import { typesContext } from "../../../../contexts/typesContext";
 import { ParameterComponentType } from "../../../../types/components";
+import { cleanEdges } from "../../../../util/reactflowUtils";
 import {
   classNames,
   getRandomKeyByssmm,
@@ -302,7 +303,13 @@ export default function ParameterComponent({
               field_name={name}
               setNodeClass={(nodeClass) => {
                 data.node = nodeClass;
-                // cleanEdges({ flow: { edges: reactFlowInstance.getEdges(), nodes: reactFlowInstance.getNodes() }, updateEdge: (edge) => reactFlowInstance.setEdges(edge) });
+                cleanEdges({
+                  flow: {
+                    edges: reactFlowInstance.getEdges(),
+                    nodes: reactFlowInstance.getNodes(),
+                  },
+                  updateEdge: (edge) => reactFlowInstance.setEdges(edge),
+                });
               }}
               nodeClass={data.node}
               disabled={disabled}
diff --git a/src/frontend/src/modals/genericModal/index.tsx b/src/frontend/src/modals/genericModal/index.tsx
index 179bef869..e3ccfce38 100644
--- a/src/frontend/src/modals/genericModal/index.tsx
+++ b/src/frontend/src/modals/genericModal/index.tsx
@@ -135,10 +135,9 @@ export default function GenericModal({
     postValidatePrompt(field_name, inputValue, nodeClass)
       .then((apiReturn) => {
         if (apiReturn.data) {
-          setNodeClass(apiReturn.data.frontend_node);
-
-          let inputVariables = apiReturn.data.input_variables;
-          if (inputVariables.length === 0) {
+          setNodeClass(apiReturn.data?.frontend_node);
+          let inputVariables = apiReturn.data.input_variables ?? [];
+          if (inputVariables && inputVariables.length === 0) {
             setIsEdit(true);
             setNoticeData({
               title: "Your template does not have any variables.",
@@ -159,6 +158,7 @@ export default function GenericModal({
         }
       })
       .catch((error) => {
+        console.log(error);
         setIsEdit(true);
         return setErrorData({
           title: "There is something wrong with this prompt, please review it",
diff --git a/src/frontend/src/util/reactflowUtils.ts b/src/frontend/src/util/reactflowUtils.ts
index e7d2dc381..5a7754ec8 100644
--- a/src/frontend/src/util/reactflowUtils.ts
+++ b/src/frontend/src/util/reactflowUtils.ts
@@ -5,6 +5,7 @@ export function cleanEdges({
   flow: { edges, nodes },
   updateEdge,
 }: cleanEdgesType) {
+  console.log(nodes, "nodes");
   let newEdges = _.cloneDeep(edges);
   edges.forEach((edge) => {
     // check if the source and target node still exists
@@ -20,8 +21,8 @@ export function cleanEdges({
       if (targetHandle) {
         const field = targetHandle.split("|")[1];
         const id =
-          (targetNode.data.node.template[field].input_types?.join(";") ??
-            targetNode.data.node.template[field].type) +
+          (targetNode.data.node.template[field]?.input_types?.join(";") ??
+            targetNode.data.node.template[field]?.type) +
           "|" +
           field +
           "|" +
@@ -42,6 +43,6 @@ export function cleanEdges({
       }
     }
   });
-  console.log("cleanEdges", newEdges);
+  console.log("clean end", newEdges);
   updateEdge(newEdges);
 }