| 201 | |
| 202 | @self.on(IterationEvents.MODEL_COMPLETED) |
| 203 | def _decollate_data(engine: Engine) -> None: |
| 204 | # replicate the scalar values to make sure all the items have batch dimension, then decollate |
| 205 | transform = Decollated(keys=None, detach=True) |
| 206 | if isinstance(engine.state.batch, (list, dict)): |
| 207 | engine.state.batch = transform(engine.state.batch) |
| 208 | if isinstance(engine.state.output, (list, dict)): |
| 209 | engine.state.output = transform(engine.state.output) |
| 210 | |
| 211 | def _register_postprocessing(self, posttrans: Callable) -> None: |
| 212 | """ |