diff --git a/trainer/train_sft_omni.py b/trainer/train_sft_omni.py
index 9df8c8e..612f5e7 100644
--- a/trainer/train_sft_omni.py
+++ b/trainer/train_sft_omni.py
@@ -4,6 +4,7 @@ import sys
 __package__ = "trainer"
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
 
+import datasets
 import argparse
 import time
 import warnings
diff --git a/trainer/trainer_utils.py b/trainer/trainer_utils.py
index 59afa41..6789e94 100644
--- a/trainer/trainer_utils.py
+++ b/trainer/trainer_utils.py
@@ -5,6 +5,7 @@ import os
 import random
 import math
 import numpy as np
+import datasets
 import torch
 import torch.distributed as dist
 from torch.utils.data import Sampler