1 files changed, 2 insertions, 5 deletions
diff --git a/lib/banzai/filter/references/iteration_reference_filter.rb b/lib/banzai/filter/references/iteration_reference_filter.rb
index cf3d446147f..591e07013c3 100644
--- a/lib/banzai/filter/references/iteration_reference_filter.rb
+++ b/lib/banzai/filter/references/iteration_reference_filter.rb
@@ -6,13 +6,10 @@ module Banzai
       # The actual filter is implemented in the EE mixin
       class IterationReferenceFilter < AbstractReferenceFilter
         self.reference_type = :iteration
-
-        def self.object_class
-          Iteration
-        end
+        self.object_class   = Iteration
       end
     end
   end
 end
 
-Banzai::Filter::References::IterationReferenceFilter.prepend_if_ee('EE::Banzai::Filter::References::IterationReferenceFilter')
+Banzai::Filter::References::IterationReferenceFilter.prepend_mod_with('Banzai::Filter::References::IterationReferenceFilter')