diff --git a/trainer/train_pretrain_vlm.py b/trainer/train_pretrain_vlm.py
index 403a1cb..1b9aca7 100644
--- a/trainer/train_pretrain_vlm.py
+++ b/trainer/train_pretrain_vlm.py
@@ -7,6 +7,7 @@ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 import argparse
 import time
 import warnings
+import datasets
 import torch
 import torch.distributed as dist
 from contextlib import nullcontext
diff --git a/trainer/train_sft_vlm.py b/trainer/train_sft_vlm.py
index 60972a8..94100ef 100644
--- a/trainer/train_sft_vlm.py
+++ b/trainer/train_sft_vlm.py
@@ -7,6 +7,7 @@ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 import argparse
 import time
 import warnings
+import datasets
 import torch
 import torch.distributed as dist
 from contextlib import nullcontext