diff --git a/toolkit/dataloader_mixins.py b/toolkit/dataloader_mixins.py
index 4b1e86d4..5a6aa239 100644
--- a/toolkit/dataloader_mixins.py
+++ b/toolkit/dataloader_mixins.py
@@ -111,12 +111,12 @@ class CaptionMixin:
         if not hasattr(self, 'file_list'):
             raise Exception('file_list not found on class instance')
         img_path_or_tuple = self.file_list[index]
+        ext = self.dataset_config.caption_ext
         if isinstance(img_path_or_tuple, tuple):
             img_path = img_path_or_tuple[0] if isinstance(img_path_or_tuple[0], str) else img_path_or_tuple[0].path
             # check if either has a prompt file
             path_no_ext = os.path.splitext(img_path)[0]
             prompt_path = None
-            ext = self.dataset_config.caption_ext
             prompt_path = path_no_ext + ext
         else:
             img_path = img_path_or_tuple if isinstance(img_path_or_tuple, str) else img_path_or_tuple.path