diff --git a/lleaves/data_processing.py b/lleaves/data_processing.py
index 30fb6e7..d8aa78c 100644
--- a/lleaves/data_processing.py
+++ b/lleaves/data_processing.py
@@ -4,7 +4,6 @@
 from typing import List, Optional
 
 import numpy as np
-import pandas as pd
 
 try:
     from pandas import DataFrame as pd_DataFrame
@@ -16,7 +15,7 @@ class pd_DataFrame:
         pass
 
 
-def _dataframe_to_ndarray(data: pd.DataFrame, pd_traintime_categories: List[List]):
+def _dataframe_to_ndarray(data: pd_DataFrame, pd_traintime_categories: List[List]):
     """
     Converts the given dataframe into a 2D numpy array and converts categorical columns to float.