ExtremeBatchCaseVisualizationCallback now has access to additional_batch_items

src/super_gradients/training/sg_trainer/sg_trainer.py

@@ -474,7 +474,9 @@ def _train_epoch(self, context: PhaseContext, silent_mode: bool = False) -> tupl
                 if self.pre_prediction_callback is not None:
                     inputs, targets = self.pre_prediction_callback(inputs, targets, batch_idx)
 
-                context.update_context(batch_idx=batch_idx, inputs=inputs, target=targets, **additional_batch_items)
+                context.update_context(
+                    batch_idx=batch_idx, inputs=inputs, target=targets, additional_batch_items=additional_batch_items, **additional_batch_items
+                )
                 self.phase_callback_handler.on_train_batch_start(context)
 
                 # AUTOCAST IS ENABLED ONLY IF self.training_params.mixed_precision - IF enabled=False AUTOCAST HAS NO EFFECT
@@ -2097,7 +2099,9 @@ def evaluate(
                     inputs, targets, additional_batch_items = sg_trainer_utils.unpack_batch_items(batch_items)
 
                     # TRIGGER PHASE CALLBACKS CORRESPONDING TO THE EVALUATION TYPE
-                    context.update_context(batch_idx=batch_idx, inputs=inputs, target=targets, **additional_batch_items)
+                    context.update_context(
+                        batch_idx=batch_idx, inputs=inputs, target=targets, additional_batch_items=additional_batch_items, **additional_batch_items
+                    )
                     if evaluation_type == EvaluationType.VALIDATION:
                         self.phase_callback_handler.on_validation_batch_start(context)
                     else:

src/super_gradients/training/utils/callbacks/base_callbacks.py

@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import List
+from typing import List, Any
 
 from typing import Optional
 import torch
@@ -63,6 +63,7 @@ def __init__(
         valid_metrics: Optional[MetricCollection] = None,  # noqa: ignore
         ema_model: Optional["SgModule"] = None,  # noqa: ignore
         loss_logging_items_names: Optional[List[str]] = None,
+        additional_batch_items: Optional[Any] = None,
     ):
         self.epoch = epoch
         self.batch_idx = batch_idx
@@ -94,6 +95,7 @@ def __init__(
         self.valid_metrics = valid_metrics
         self.ema_model = ema_model
         self.loss_logging_items_names = loss_logging_items_names
+        self.additional_batch_items = additional_batch_items
 
     def update_context(self, **kwargs):
         for attr, attr_val in kwargs.items():

src/super_gradients/training/utils/callbacks/callbacks.py

@@ -1131,6 +1131,7 @@ def __init__(
         self.extreme_batch = None
         self.extreme_preds = None
         self.extreme_targets = None
+        self.extreme_additional_batch_items = None
 
         self._first_call = True
         self._idx_loss_tuple = None
@@ -1224,6 +1225,7 @@ def _on_batch_end(self, context: PhaseContext) -> None:
             self.extreme_batch = tensor_container_to_device(context.inputs, device="cpu", detach=True, non_blocking=False)
             self.extreme_preds = tensor_container_to_device(context.preds, device="cpu", detach=True, non_blocking=False)
             self.extreme_targets = tensor_container_to_device(context.target, device="cpu", detach=True, non_blocking=False)
+            self.extreme_additional_batch_items = tensor_container_to_device(context.additional_batch_items, device="cpu", detach=True, non_blocking=False)
 
     def _init_loss_attributes(self, context: PhaseContext):
         if self.loss_to_monitor not in context.loss_logging_items_names:
@@ -1236,6 +1238,7 @@ def _reset(self):
         self.extreme_batch = None
         self.extreme_preds = None
         self.extreme_targets = None
+        self.extreme_additional_batch_items = None
         if self.metric is not None:
             self.metric.reset()