diff --git a/src/mflux/community/in_context_lora/flux_in_context_lora.py b/src/mflux/community/in_context_lora/flux_in_context_lora.py
index 3a823fc..9c27fc2 100644
--- a/src/mflux/community/in_context_lora/flux_in_context_lora.py
+++ b/src/mflux/community/in_context_lora/flux_in_context_lora.py
@@ -114,7 +114,7 @@ class Flux1InContextLoRA(nn.Module):
                     static_noise=static_noise,
                 )
 
-                # (Optional) Call subscribers at end of loop
+                # (Optional) Call subscribers in-loop
                 Callbacks.in_loop(
                     t=t,
                     seed=seed,
@@ -137,6 +137,14 @@ class Flux1InContextLoRA(nn.Module):
                     time_steps=time_steps,
                 )
 
+        # (Optional) Call subscribers after loop
+        Callbacks.after_loop(
+            seed=seed,
+            prompt=prompt,
+            latents=latents,
+            config=config,
+        )  # fmt: off
+
         # 6. Decode the latent array and return the image
         latents = ArrayUtil.unpack_latents(latents=latents, height=config.height, width=config.width)
         decoded = self.vae.decode(latents)
diff --git a/src/mflux/controlnet/flux_controlnet.py b/src/mflux/controlnet/flux_controlnet.py
index 27c4230..5c9d878 100644
--- a/src/mflux/controlnet/flux_controlnet.py
+++ b/src/mflux/controlnet/flux_controlnet.py
@@ -119,7 +119,7 @@ class Flux1Controlnet(nn.Module):
                 dt = config.sigmas[t + 1] - config.sigmas[t]
                 latents += noise * dt
 
-                # (Optional) Call subscribers at end of loop
+                # (Optional) Call subscribers in-loop
                 Callbacks.in_loop(
                     t=t,
                     seed=seed,
@@ -142,7 +142,7 @@ class Flux1Controlnet(nn.Module):
                     time_steps=time_steps,
                 )
 
-        # (Optional) Call subscribers at end of loop
+        # (Optional) Call subscribers after loop
         Callbacks.after_loop(
             seed=seed,
             prompt=prompt,
diff --git a/src/mflux/flux/flux.py b/src/mflux/flux/flux.py
index 784c4d8..bf7f756 100644
--- a/src/mflux/flux/flux.py
+++ b/src/mflux/flux/flux.py
@@ -99,7 +99,7 @@ class Flux1(nn.Module):
                 dt = config.sigmas[t + 1] - config.sigmas[t]
                 latents += noise * dt
 
-                # (Optional) Call subscribers at end of loop
+                # (Optional) Call subscribers in-loop
                 Callbacks.in_loop(
                     t=t,
                     seed=seed,
@@ -122,7 +122,7 @@ class Flux1(nn.Module):
                     time_steps=time_steps,
                 )
 
-        # (Optional) Call subscribers at end of loop
+        # (Optional) Call subscribers after loop
         Callbacks.after_loop(
             seed=seed,
             prompt=prompt,