1 files changed, 3 insertions, 6 deletions
diff --git a/Scripts/Fold/Editor/FoldGraphView.cs b/Scripts/Fold/Editor/FoldGraphView.cs
index 2a94b9f..f77e97b 100644
--- a/Scripts/Fold/Editor/FoldGraphView.cs
+++ b/Scripts/Fold/Editor/FoldGraphView.cs
@@ -3,6 +3,7 @@ using UnityEditor;
 using GraphProcessor;
 using System;
 using System.Collections.Generic;
+using System.Linq;
 
 public class FoldGraphView : BaseGraphView
 {
@@ -10,12 +11,8 @@ public class FoldGraphView : BaseGraphView
 
     public override IEnumerable<(string path, Type type)> FilterCreateNodeMenuEntries()
     {
-        // Only return our subset of nodes
-        foreach (var nodeMenuItem in NodeProvider.GetNodeMenuEntries(graph))
-        {
-            if (typeof(BaseFoldNode).IsAssignableFrom(nodeMenuItem.type))
-                yield return nodeMenuItem;
-        }
+        return NodeProvider.GetNodeMenuEntries(graph)
+            .Where(entry => typeof(BaseFoldNode).IsAssignableFrom(entry.type));
     }
 }