diff --git a/lib/liquid/parse_tree_visitor.rb b/lib/liquid/parse_tree_visitor.rb
index 62f6c7b..74f5563 100644
--- a/lib/liquid/parse_tree_visitor.rb
+++ b/lib/liquid/parse_tree_visitor.rb
@@ -17,8 +17,8 @@ module Liquid
 
     def add_callback_for(*classes, &block)
       callback = block
-      callback = ->(node, _) { block.call(node) } if block.arity.abs == 1
-      callback = ->(_, _) { block.call } if block.arity.zero?
+      callback = ->(node, _) { yield node } if block.arity.abs == 1
+      callback = ->(_, _) { yield } if block.arity.zero?
       classes.each { |klass| @callbacks[klass] = callback }
       self
     end