Diff Coverage

    def __delitem__(self, key: str) -> None:
        """Delete every registration for ``key`` from this registry."""
        if key not in self._local_mapping and key not in self._global_mapping:
            raise KeyError(key)
        self._local_mapping.pop(key, None)
        self._global_mapping.pop(key, None)

    def __iter__(self) -> Iterator[str]:

        try:
            empty_cache()
            dist.barrier()
            synchronize()
        finally:
            destroy_process_group()

    def train(self) -> None:

            start_epoch = self.state.epoch
            for epoch in range(start_epoch, self.train_epochs):
                if epoch != start_epoch:
                    self.train_dataloader.set_epoch(epoch)
                    self.data_iterator = HyperIter(
                        self.train_dataloader, use_background_prefetcher=config.dataloader.use_background_prefetcher
                    )
                self.state.epoch = epoch

                for _ in range(start_step, train_steps):
                    try:
                        self.train_step(self.data_iterator)
                    except StopIteration:
                        logger.info(
                            "epoch:%s Dataloader finished with drop_last %s",
                            epoch,
                            config.dataloader.drop_last,
                        )
                        break

                self.on_epoch_end()
                self.state.epoch = epoch + 1

                print_device_mem_info(f"VRAM usage after epoch {epoch + 1}")

                self.data_iterator.stop()
                self.data_iterator = None
        finally:
            # ExitStack runs every callback even when an earlier cleanup raises.
            with ExitStack() as cleanup:
                cleanup.callback(self.destroy_distributed)