diff --git a/src/moscot/base/problems/_mixins.py b/src/moscot/base/problems/_mixins.py
index a5d2cf0a..d69dcfdd 100644
--- a/src/moscot/base/problems/_mixins.py
+++ b/src/moscot/base/problems/_mixins.py
@@ -285,7 +285,7 @@ def _annotation_mapping(
                         key_added=None,
                     )
                     v = np.array(tm_batch.argmax(0))
-                    out.extend(source_df[annotation_label][v[i]] for i in range(len(v)))
+                    out.extend(source_df[annotation_label].iloc[v[i]] for i in range(len(v)))
 
             else:
                 target_df = _get_df_cell_transition(
@@ -309,7 +309,7 @@ def _annotation_mapping(
                         key_added=None,
                     )
                     v = np.array(tm_batch.argmax(0))
-                    out.extend(target_df[annotation_label][v[i]] for i in range(len(v)))
+                    out.extend(target_df[annotation_label].iloc[v[i]] for i in range(len(v)))
             categories = pd.Categorical(out)
             return pd.DataFrame(categories, columns=[annotation_label])
         raise NotImplementedError(f"Mapping mode `{mapping_mode!r}` is not yet implemented.")
diff --git a/tests/problems/base/test_general_problem.py b/tests/problems/base/test_general_problem.py
index 7cd907ba..8a4abee6 100644
--- a/tests/problems/base/test_general_problem.py
+++ b/tests/problems/base/test_general_problem.py
@@ -9,6 +9,7 @@
 from ott.geometry.pointcloud import PointCloud
 from ott.solvers.linear import solve as sinkhorn
 
+import anndata as ad
 import scanpy as sc
 from anndata import AnnData
 
@@ -205,7 +206,7 @@ def test_set_graph_xy(self, adata_x: AnnData, adata_y: AnnData, ts: Tuple[Option
         new_obs_names = [name + "_src" for name in adata_x.obs_names]
         adata_x.obs_names = new_obs_names
 
-        adata_concat = adata_x.concatenate(adata_y, index_unique=None)
+        adata_concat = ad.concat([adata_x, adata_y], index_unique=None)
         sc.pp.neighbors(adata_concat, n_neighbors=15)
         graph_to_set = pd.DataFrame(
             index=adata_concat.obs_names,
@@ -323,7 +324,7 @@ def test_set_graph_xy_test_t(self, adata_x: AnnData, adata_y: AnnData, t: float)
         new_obs_names = [name + "_src" for name in adata_x.obs_names]
         adata_x.obs_names = new_obs_names
 
-        adata_concat = adata_x.concatenate(adata_y, index_unique=None)
+        adata_concat = ad.concat([adata_x, adata_y], index_unique=None)
         sc.pp.neighbors(adata_concat, n_neighbors=5)
         graph_to_set = pd.DataFrame(
             index=adata_concat.obs_names,