diff --git a/extensions-builtin/forge_preprocessor_reference/scripts/forge_reference.py b/extensions-builtin/forge_preprocessor_reference/scripts/forge_reference.py
index bb38ceef..0461ea1b 100644
--- a/extensions-builtin/forge_preprocessor_reference/scripts/forge_reference.py
+++ b/extensions-builtin/forge_preprocessor_reference/scripts/forge_reference.py
@@ -70,7 +70,9 @@ class PreprocessorReference(Preprocessor):
                 return h
 
             if flag != 'after':
-                return
+                return h
+
+            location = transformer_options['block']
 
             sigma = transformer_options["sigmas"][0].item()
             if not (sigma_min <= sigma <= sigma_max):
@@ -91,6 +93,9 @@ class PreprocessorReference(Preprocessor):
             if not (sigma_min <= sigma <= sigma_max):
                 return sdp(q, k, v, transformer_options)
 
+            location = (transformer_options['block'][0], transformer_options['block'][1],
+                        transformer_options['block_index'])
+
             if self.is_recording_style:
                 a = 0
             else: