def teardown(self) -> None: super().teardown() if isinstance(self.model, DistributedDataParallel): self.model = self.lightning_module if self.sync_batchnorm: self.model = _revert_sync_batchnorm(self.model) if self.root_device.type == "cuda": # GPU teardown self.lightning_module.cpu() # clean up memory torch.cuda.empty_cache()
def teardown(self) -> None: log.detail(f"{self.__class__.__name__}: tearing down DDP plugin") super().teardown() if isinstance(self.model, DistributedDataParallel): self.model = self.lightning_module if self.sync_batchnorm: self.model = _revert_sync_batchnorm(self.model) if self.root_device.type == "cuda": # GPU teardown log.detail(f"{self.__class__.__name__}: moving model to CPU") self.lightning_module.cpu() # clean up memory torch.cuda.empty_cache()