Check if iterationIndex is nil

Signed-off-by: droctothorpe <[email protected]>
Co-authored-by: zazulam <[email protected]>

backend/src/v2/driver/driver.go

@@ -1245,8 +1245,11 @@ func getDAGTasks(
 			// within a single DAG, which results in an error when
 			// mlmd.GetExecutionsInDAG is called. ParallelFor outputs should be
 			// handled with dsl.Collected.
-			if _, ok := v.GetExecution().GetCustomProperties()["iteration_index"]; !ok {
-				continue
+			if iterationIndex, ok := v.GetExecution().GetCustomProperties()["iteration_index"]; !ok {
+				// First we check if the key exists, then we check if it's not nil.
+				if !(iterationIndex == nil) {
+					continue
+				}
 			}
 			glog.V(4).Infof("Found a task, %v, with an execution type of system.DAGExecution. Adding its tasks to the task list.", v.TaskName())
 			subDAG, err := mlmd.GetDAG(ctx, v.GetExecution().GetId())