Add gradient accumulation supervised trainer update step function

Signed-off-by: Jakob Weigand <[email protected]>

monai/engines/utils.py

@@ -18,7 +18,7 @@
 import torch
 
 from monai.config import IgniteInfo
-from monai.engines import Trainer
+from monai.engines import SupervisedTrainer
 from monai.transforms import apply_transform
 from monai.utils import ensure_tuple, min_version, optional_import
 from monai.utils.enums import CommonKeys, GanKeys
@@ -306,7 +306,7 @@ def __init__(self, gradient_accumulation_steps: int = 1) -> None:
                             "No gradient accumulation if the value set to one (default).")
         self.gradient_accumulation_steps = gradient_accumulation_steps
 
-    def __call__(self, engine: Trainer, batchdata: Sequence[torch.Tensor]) -> Any | tuple[torch.Tensor]:
+    def __call__(self, engine: SupervisedTrainer, batchdata: Sequence[torch.Tensor]) -> Any | tuple[torch.Tensor]:
         if batchdata is None:
             raise ValueError("Must provide batch data for current iteration.")
         batch = engine.prepare_batch(batchdata, engine.state.device, engine.non_blocking, **engine.to_kwargs)