From 596868e4ca2933050810706d534c796a2592d010 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Pstr=C4=85g?= <michal.pstrag@icloud.com>
Date: Sat, 17 Aug 2024 03:05:08 +0200
Subject: [PATCH] fix types issues

---
 extra/prompt_tuning/tuning/loaders.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/extra/prompt_tuning/tuning/loaders.py b/extra/prompt_tuning/tuning/loaders.py
index 2cc7dc96..9553643c 100644
--- a/extra/prompt_tuning/tuning/loaders.py
+++ b/extra/prompt_tuning/tuning/loaders.py
@@ -1,8 +1,9 @@
 from abc import ABC, abstractmethod
-from typing import Dict, Iterable, List
+from typing import Iterable, List
 
 import dspy.datasets
 from dspy import Example
+from omegaconf import DictConfig
 
 
 class DataLoader(ABC):
@@ -10,7 +11,7 @@ class DataLoader(ABC):
     Data loader.
     """
 
-    def __init__(self, config: Dict) -> None:
+    def __init__(self, config: DictConfig) -> None:
         self.config = config
 
     @abstractmethod