Merge pull request #180 from OFA-Sys/feature/logicwong

Feature/logicwong
OFA-Sys · Jul 25, 2022 · 9c1a516 · 9c1a516
2 parents c8313b9 + b200f52
commit 9c1a516
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tasks/pretrain_tasks/unify_task.py b/tasks/pretrain_tasks/unify_task.py
@@ -7,8 +7,10 @@
 import json
 import logging
 import os
+import math
 from typing import Optional
 from fairseq.tasks import register_task
+from fairseq.data import FairseqDataset, iterators
 
 from tasks.ofa_task import OFATask, OFAConfig
 from data.pretrain_data.unify_dataset import UnifyDataset