diff --git a/src/graphnet/data/dataset/parquet/parquet_dataset.py b/src/graphnet/data/dataset/parquet/parquet_dataset.py
index 3561c591a..40fc9805b 100644
--- a/src/graphnet/data/dataset/parquet/parquet_dataset.py
+++ b/src/graphnet/data/dataset/parquet/parquet_dataset.py
@@ -5,6 +5,7 @@
     List,
     Optional,
     Union,
+    Any,
 )
 
 import numpy as np