revise the package import error in tools/caption/caption_llava.py (#357)

hpcaitech · May 9, 2024 · c3fd95f · c3fd95f
1 parent 3e4e8a4
commit c3fd95f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tools/caption/caption_llava.py b/tools/caption/caption_llava.py
@@ -18,7 +18,8 @@
 from tqdm import tqdm
 
 from .acceleration.llava.policies import LlavaLlamaForCausalLMPolicy, LlavaMistralForCausalLMPolicy
-from .utils import IMG_EXTENSIONS, PROMPTS, VID_EXTENSIONS, Timer, VideoTextDataset, collate_fn
+from .utils import PROMPTS, Timer, VideoTextDataset, collate_fn
+from ..datasets.utils import IMG_EXTENSIONS, VID_EXTENSIONS
 
 disable_torch_init()