diff --git a/servers/visual/shader_language.cpp b/servers/visual/shader_language.cpp
index 559edc2282..70863dc86c 100644
--- a/servers/visual/shader_language.cpp
+++ b/servers/visual/shader_language.cpp
@@ -5385,11 +5385,6 @@ Error ShaderLanguage::_parse_block(BlockNode *p_block, const Map<StringName, Bui
 			//check return type
 			BlockNode *b = p_block;
 
-			if (b && b->parent_function && (b->parent_function->name == "vertex" || b->parent_function->name == "fragment" || b->parent_function->name == "light")) {
-				_set_error(vformat("Using 'return' in '%s' processor function results in undefined behavior!", b->parent_function->name));
-				return ERR_PARSE_ERROR;
-			}
-
 			while (b && !b->parent_function) {
 				b = b->parent_block;
 			}
@@ -5399,6 +5394,11 @@ Error ShaderLanguage::_parse_block(BlockNode *p_block, const Map<StringName, Bui
 				return ERR_BUG;
 			}
 
+			if (b && b->parent_function && (b->parent_function->name == "vertex" || b->parent_function->name == "fragment" || b->parent_function->name == "light")) {
+				_set_error(vformat("Using 'return' in '%s' processor function results in undefined behavior!", b->parent_function->name));
+				return ERR_PARSE_ERROR;
+			}
+
 			String return_struct_name = String(b->parent_function->return_struct_name);
 
 			ControlFlowNode *flow = alloc_node<ControlFlowNode>();