added if statement

AIRI-Institute · Mar 23, 2024 · 3a184b1 · 3a184b1
1 parent 20a276c
commit 3a184b1
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/probing/pipeline.py b/probing/pipeline.py
@@ -1,6 +1,6 @@
 import os
 from time import time
-from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union, get_args
 
 import numpy as np
 import torch
@@ -170,7 +170,7 @@ def run(
         save_checkpoints: bool = False,
         verbose: bool = True,
     ) -> None:
-        if path_to_task_file:
+        if path_to_task_file or probe_task in get_args(UDProbingTaskName):
             task_data = TextFormer(probe_task, path_to_task_file)
             task_dataset, num_classes = task_data.samples, len(task_data.unique_labels)
             probing_task_language, probing_task_category = lang_category_extraction(
@@ -191,7 +191,7 @@ def run(
 
         clear_memory()
         start_time = time()
-        if path_to_task_file:
+        if path_to_task_file or probe_task in get_args(UDProbingTaskName):
             (
                 probing_dataloaders,
                 mapped_labels,