update example

examples/callback/custom_callback.py

@@ -11,13 +11,17 @@
 from benchmarl.experiment import Experiment, ExperimentConfig
 from benchmarl.experiment.callback import Callback
 from benchmarl.models.mlp import MlpConfig
-from tensordict import TensorDictBase
+from tensordict import TensorDict, TensorDictBase
 
 
 class MyCallbackA(Callback):
     def on_batch_collected(self, batch: TensorDictBase):
         print(f"Callback A is doing something with the sampling batch {batch}")
 
+    def on_train_step(self, batch: TensorDictBase, group: str) -> TensorDictBase:
+        print(f"Callback A is computing a loss with the training tensordict {batch}")
+        return TensorDict({}, [])
+
     def on_train_end(self, training_td: TensorDictBase, group: str):
         print(
             f"Callback A is doing something with the training tensordict {training_td}"