diff --git a/src/renderer/components/Experiment/Workflows/NewNodeModal.tsx b/src/renderer/components/Experiment/Workflows/NewNodeModal.tsx
index d9fe123a516317bf806c49fda59915f509cdb969..30cf6045cbff57e380161f354cb8c0ed6ac8de91 100644
--- a/src/renderer/components/Experiment/Workflows/NewNodeModal.tsx
+++ b/src/renderer/components/Experiment/Workflows/NewNodeModal.tsx
@@ -62,7 +62,6 @@ export default function NewNodeModal({
               const node = {
                 name: name,
                 type: 'TRAIN',
-                out: (config.nodes.length + 1).toString(),
                 template: template,
               };
               await fetch(
@@ -78,7 +77,6 @@ export default function NewNodeModal({
               const node = {
                 name: name,
                 type: 'EVAL',
-                out: (config.nodes.length + 1).toString(),
                 template: template,
               };
               await fetch(
diff --git a/src/renderer/components/Experiment/Workflows/WorkflowCanvas.tsx b/src/renderer/components/Experiment/Workflows/WorkflowCanvas.tsx
index 41c3206d12a1277a8b8df3a1c4b089af6ce50f8b..0f3c3b79e91c4a5abf129528af1bdb08544c00b3 100644
--- a/src/renderer/components/Experiment/Workflows/WorkflowCanvas.tsx
+++ b/src/renderer/components/Experiment/Workflows/WorkflowCanvas.tsx
@@ -14,18 +14,30 @@ import CustomNode from './CustomNode';
 const nodeTypes = { customNode: CustomNode };
 
 function generateNodes(workflow: any) {
+  const workflowConfig = JSON.parse(workflow?.config);
+
+  if (workflowConfig.nodes.length == 0) {
+    return [];
+  }
+
   let out: any[] = [];
-  let currentTask = '0';
+  let currentTask = workflowConfig.nodes[0].id;
   let position = 0;
 
-  const workflowConfig = JSON.parse(workflow?.config);
   console.log(workflowConfig);
 
-  while (currentTask < workflowConfig.nodes.length) {
+  while (currentTask != 'END') {
+    let currentNode = {};
+    workflowConfig.nodes.forEach((node) => {
+      if (node.id == currentTask) {
+        currentNode = node;
+      }
+    });
+
     const data = {
-      label: workflowConfig.nodes[currentTask].name,
-      jobType: workflowConfig.nodes[currentTask].type,
-      template: workflowConfig.nodes[currentTask].template,
+      label: currentNode.name,
+      jobType: currentNode.type,
+      template: currentNode.template,
     };
     const nextNode = {
       id: currentTask,
@@ -35,31 +47,42 @@ function generateNodes(workflow: any) {
     };
     out.push(nextNode);
     position += 120;
-    currentTask = workflowConfig.nodes[currentTask].out;
+    currentTask = currentNode.out;
   }
 
   return out;
 }
 
 function generateEdges(workflow: any) {
+  const workflowConfig = JSON.parse(workflow?.config);
+  if (workflowConfig.nodes.length <= 1) {
+    return [];
+  }
+
   let out: any[] = [];
-  let currentTask = '0';
-  let ids = '0';
+  let currentTask = workflowConfig.nodes[0].id;
+  let ids = workflowConfig.nodes[0].id;
 
-  const workflowConfig = JSON.parse(workflow?.config);
   console.log(workflowConfig);
 
-  while (currentTask < workflowConfig.nodes.length) {
+  while (currentTask != 'END') {
+    let currentNode = {};
+    workflowConfig.nodes.forEach((node) => {
+      if (node.id == currentTask) {
+        currentNode = node;
+      }
+    });
+
     out.push({
       id: ids,
       source: currentTask,
-      target: workflowConfig.nodes[currentTask].out,
+      target: currentNode.out,
       markerEnd: {
         type: 'arrow',
       },
     });
     ids += 1;
-    currentTask = workflowConfig.nodes[currentTask].out;
+    currentTask = currentNode.out;
   }
 
   return out;