fix bugs in unify_task.py

OFA-Sys · Jul 25, 2022 · 05975e6 · 05975e6
1 parent 3b4952a
commit 05975e6
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