From bbc19ec084db023ca47ea40c3b3e9787610cc7a1 Mon Sep 17 00:00:00 2001
From: Harshita Sharma <61745616+harshita-s@users.noreply.github.com>
Date: Thu, 13 Oct 2022 11:27:39 +0100
Subject: [PATCH] BUG: Checkpoint path not picked up correctly (#630)

Fix for the checkpoint path (of the best metrics checkpoint) not being
correctly picked up on the cluster.
---
 .../src/health_cpath/configs/classification/BaseMIL.py        | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/hi-ml-cpath/src/health_cpath/configs/classification/BaseMIL.py b/hi-ml-cpath/src/health_cpath/configs/classification/BaseMIL.py
index 829e8b0d..edb646e3 100644
--- a/hi-ml-cpath/src/health_cpath/configs/classification/BaseMIL.py
+++ b/hi-ml-cpath/src/health_cpath/configs/classification/BaseMIL.py
@@ -196,6 +196,10 @@ class BaseMIL(LightningContainer, EncoderParams, PoolingParams, LossCallbackPara
         if absolute_checkpoint_path_parent.is_file():
             return absolute_checkpoint_path_parent
 
+        checkpoint_path = Path(self.checkpoint_folder, self.best_checkpoint_filename_with_suffix)
+        if checkpoint_path.is_file():
+            return checkpoint_path
+
         checkpoint_path = get_best_checkpoint_path(self.checkpoint_folder)
         if checkpoint_path.is_file():
             return checkpoint_path