diff --git a/numalogic/udfs/trainer/_base.py b/numalogic/udfs/trainer/_base.py
index 62ae7a63..24a48524 100644
--- a/numalogic/udfs/trainer/_base.py
+++ b/numalogic/udfs/trainer/_base.py
@@ -256,6 +256,7 @@ def _is_data_sufficient(self, payload: TrainerPayload, df: pd.DataFrame) -> bool
             return False
         return True
 
+    # TODO: Use a custom imputer in transforms module
     @staticmethod
     def get_feature_arr(
         raw_df: pd.DataFrame, metrics: list[str], fill_value: float = 0.0
@@ -265,7 +266,7 @@ def get_feature_arr(
             if col not in raw_df.columns:
                 raw_df[col] = fill_value
         feat_df = raw_df[metrics]
-        feat_df = feat_df.fillna(fill_value)
+        feat_df = feat_df.fillna(fill_value).replace([np.inf, -np.inf], fill_value)
         return feat_df.to_numpy(dtype=np.float32)
 
     def fetch_data(self, payload: TrainerPayload) -> pd.DataFrame:
diff --git a/pyproject.toml b/pyproject.toml
index 662faf06..d2cb3795 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "numalogic"
-version = "0.6.1.dev3"
+version = "0.6.1.dev4"
 description = "Collection of operational Machine Learning models and tools."
 authors = ["Numalogic Developers"]
 packages = [{ include = "numalogic" }]