diff --git a/declearn/dataset/_inmemory.py b/declearn/dataset/_inmemory.py
index a467de216d43e68f09451f3659ae59c93d0cf81b..58bd35f20b194c3d9fb62c952db0c17f4d7ace13 100644
--- a/declearn/dataset/_inmemory.py
+++ b/declearn/dataset/_inmemory.py
@@ -177,8 +177,8 @@ class InMemoryDataset(Dataset):
             return self.data
         if isinstance(self.data, pd.DataFrame):
             if isinstance(self.f_cols[-1], str):
-                return self.data.loc[:, self.f_cols]
-            return self.data.iloc[:, self.f_cols]
+                return self.data.loc[:, self.f_cols]  # type: ignore
+            return self.data.iloc[:, self.f_cols]  # type: ignore
         return self.data[:, self.f_cols]  # type: ignore
 
     @property
@@ -187,7 +187,7 @@ class InMemoryDataset(Dataset):
         if (not self.expose_classes) or (self.target is None):
             return None
         if isinstance(self.target, pd.DataFrame):
-            return set(self.target.unstack().unique().tolist())
+            return set(self.target.unstack().unique().tolist())  # type: ignore
         if isinstance(self.target, pd.Series):
             return set(self.target.unique().tolist())
         if isinstance(self.target, np.ndarray):