diff --git a/nerfactory/engine/trainer.py b/nerfactory/engine/trainer.py
index ff9c496f8..bd5216c5d 100644
--- a/nerfactory/engine/trainer.py
+++ b/nerfactory/engine/trainer.py
@@ -138,6 +138,11 @@ class Trainer:
                         )
 
                     loss_metric_dict = self.train_iteration(step)
+
+                    # training callbacks after the training iteration
+                    for callback in self.callbacks:
+                        callback.run_callback_at_location(step, location=TrainingCallbackLocation.AFTER_TRAIN_ITERATION)
+
                     with TimeWriter(writer, EventName.ITER_VIS_TIME, step=step) as _:
                         self.visualizer_state.update_scene(step, self.pipeline.model)
 
diff --git a/nerfactory/models/vanilla_nerf.py b/nerfactory/models/vanilla_nerf.py
index 70ba6ce46..0f5dd1548 100644
--- a/nerfactory/models/vanilla_nerf.py
+++ b/nerfactory/models/vanilla_nerf.py
@@ -85,7 +85,7 @@ class NeRFModel(Model):
                     where_to_run=[TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
                     update_every_num_iters=self.density_field.update_every_num_iters,
                     func=self.density_field.update_density_grid,
-                    kwargs={"density_eval_func": self.field.density_fn},  # type: ignore
+                    kwargs={"density_eval_func": self.field_coarse.density_fn},  # type: ignore
                 )
             ]
         return callbacks