import

microsoft · Aug 12, 2024 · 83ceef2 · 83ceef2
1 parent 245ac31
commit 83ceef2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mttl/datamodule/mmlu_data_module.py b/mttl/datamodule/mmlu_data_module.py
@@ -192,7 +192,7 @@ class MMLUDataModule(DataModule):
 
     def test_dataloader(self, subsample=None, shuffle=False):
         if subsample is not None and subsample > 0:
-            from mttl.datamodule import take_n_examples_per_task
+            from mttl.datamodule.base import take_n_examples_per_task
 
             indices = take_n_examples_per_task(
                 list(self.test_dataset["Task"]),