diff --git a/pipeline/mimicit_utils/mimicit_dataset.py b/pipeline/mimicit_utils/mimicit_dataset.py
index 3213b500..8c148767 100755
--- a/pipeline/mimicit_utils/mimicit_dataset.py
+++ b/pipeline/mimicit_utils/mimicit_dataset.py
@@ -12,6 +12,7 @@
 from io import BytesIO
 import pandas as pd
 import numpy as np
+import pyarrow.parquet as pq
 import orjson
 import torch
 from PIL import Image, ImageFile