diff --git a/src/axolotl/evaluate.py b/src/axolotl/evaluate.py index f3be9c2f4..216ff0110 100644 --- a/src/axolotl/evaluate.py +++ b/src/axolotl/evaluate.py @@ -15,6 +15,7 @@ from axolotl.logging_config import configure_logging from axolotl.train import TrainDatasetMeta from axolotl.utils import set_pytorch_cuda_alloc_conf from axolotl.utils.dict import DictDefault +from axolotl.utils.distributed import cleanup_distributed from axolotl.utils.models import load_model, load_processor, load_tokenizer from axolotl.utils.trainer import setup_trainer @@ -159,4 +160,6 @@ def evaluate(*, cfg: DictDefault, dataset_meta: TrainDatasetMeta) -> Dict[str, f del model del tokenizer + cleanup_distributed() + return all_metrics