diff --git a/src/TensorFlowNET.Core/Data/DatasetV2.cs b/src/TensorFlowNET.Core/Data/DatasetV2.cs
index 00b04edf..763baa31 100644
--- a/src/TensorFlowNET.Core/Data/DatasetV2.cs
+++ b/src/TensorFlowNET.Core/Data/DatasetV2.cs
@@ -132,7 +132,9 @@ namespace Tensorflow
                     break;
                 }
 
-                yield return (results[0], results.Length == 1 ? null : results[1]);
+                yield return results.Length == 2
+                    ? (results[0], results[1])
+                    : (null, results[0]);
             }
         }