diff --git a/python/ray/air/util/data_batch_conversion.py b/python/ray/air/util/data_batch_conversion.py
index 426814d72f1d..e134b5b1d31f 100644
--- a/python/ray/air/util/data_batch_conversion.py
+++ b/python/ray/air/util/data_batch_conversion.py
@@ -319,7 +319,7 @@ def _cast_ndarray_columns_to_tensor_extension(df: "pd.DataFrame") -> "pd.DataFra
                 with warnings.catch_warnings():
                     warnings.simplefilter("ignore", category=FutureWarning)
                     warnings.simplefilter("ignore", category=SettingWithCopyWarning)
-                    df.loc[:, col_name] = TensorArray(col)
+                    df[col_name] = TensorArray(col)
             except Exception as e:
                 raise ValueError(
                     f"Tried to cast column {col_name} to the TensorArray tensor "
@@ -354,5 +354,5 @@ def _cast_tensor_columns_to_ndarrays(df: "pd.DataFrame") -> "pd.DataFrame":
             with warnings.catch_warnings():
                 warnings.simplefilter("ignore", category=FutureWarning)
                 warnings.simplefilter("ignore", category=SettingWithCopyWarning)
-                df.loc[:, col_name] = pd.Series(list(col.to_numpy()))
+                df[col_name] = list(col.to_numpy())
     return df
diff --git a/python/ray/data/tests/test_tensor.py b/python/ray/data/tests/test_tensor.py
index 852142bf4150..3932d5b76ac0 100644
--- a/python/ray/data/tests/test_tensor.py
+++ b/python/ray/data/tests/test_tensor.py
@@ -564,7 +564,7 @@ def test_tensors_in_tables_pandas_roundtrip(
     ds_df = ds.to_pandas()
     expected_df = df + 1
     if enable_automatic_tensor_extension_cast:
-        expected_df.loc[:, "two"] = list(expected_df["two"].to_numpy())
+        expected_df["two"] = list(expected_df["two"].to_numpy())
     pd.testing.assert_frame_equal(ds_df, expected_df)
 
 
@@ -585,7 +585,7 @@ def test_tensors_in_tables_pandas_roundtrip_variable_shaped(
     ds_df = ds.to_pandas()
     expected_df = df + 1
     if enable_automatic_tensor_extension_cast:
-        expected_df.loc[:, "two"] = _create_possibly_ragged_ndarray(
+        expected_df["two"] = _create_possibly_ragged_ndarray(
             expected_df["two"].to_numpy()
         )
     pd.testing.assert_frame_equal(ds_df, expected_df)
@@ -873,8 +873,8 @@ def test_tensors_in_tables_iter_batches(
     )
     df = pd.concat([df1, df2], ignore_index=True)
     if enable_automatic_tensor_extension_cast:
-        df.loc[:, "one"] = list(df["one"].to_numpy())
-        df.loc[:, "two"] = list(df["two"].to_numpy())
+        df["one"] = list(df["one"].to_numpy())
+        df["two"] = list(df["two"].to_numpy())
     ds = ray.data.from_pandas([df1, df2])
     batches = list(ds.iter_batches(batch_size=2, batch_format="pandas"))
     assert len(batches) == 3