def on_train_epoch_end(self) -> None:
        """Ensure memory bank is fitted after training.

        This hook automatically fits the memory bank if it hasn't been fitted yet.
        """
        if not self._is_fitted:
            self.fit()
            self._is_fitted = torch.tensor([True], device=self.device)
Ensure memory bank is fitted after training.

        This hook automatically fits the memory bank if it hasn't been fitted yet.
        
src/anomalib/models/components/base/memory_bank_module.py