diff --git a/docs/pages/examples/RowSelectionExample.tsx b/docs/pages/examples/RowSelectionExample.tsx
index 4829ca8..5c72ca7 100644
--- a/docs/pages/examples/RowSelectionExample.tsx
+++ b/docs/pages/examples/RowSelectionExample.tsx
@@ -1,3 +1,4 @@
+import { Button } from '@mantine/core';
 import { RowSelectionState } from '@tanstack/react-table';
 import { useState } from 'react';
 import {
@@ -11,55 +12,66 @@ import {
 import { demoData } from '../../demoData';
 
 export default function RowSelectionExample() {
+  const [data, setData] = useState(demoData);
   const [selected, setSelected] = useState<RowSelectionState>({});
+
+  const removeSelected = () => {
+    const next = data.filter((_, i) => !selected[i]);
+    setSelected({});
+    setData(next);
+  };
+
   return (
-    <DataGrid
-      data={demoData}
-      striped
-      highlightOnHover
-      withGlobalFilter
-      withPagination
-      withColumnFilters
-      withSorting
-      withColumnResizing
-      withRowSelection
-      state={{
-        rowSelection: selected,
-      }}
-      onRowSelectionChange={setSelected}
-      columns={[
-        {
-          accessorKey: 'text',
-          header: 'Text that is too long for a Header',
-          filterFn: stringFilterFn,
-          size: 300,
-          cell: highlightFilterValue,
-        },
-        {
-          header: 'Animal',
-          columns: [
-            { accessorKey: 'cat', filterFn: stringFilterFn },
-            {
-              accessorKey: 'fish',
-              filterFn: stringFilterFn,
-            },
-          ],
-        },
-        {
-          accessorKey: 'city',
-          filterFn: stringFilterFn,
-        },
-        { accessorKey: 'value', filterFn: numberFilterFn },
-        {
-          accessorKey: 'date',
-          cell: (cell) => cell.getValue<Date>()?.toLocaleDateString(),
-          filterFn: dateFilterFn,
-        },
-        {
-          accessorKey: 'bool',
-          filterFn: booleanFilterFn,
-        },
-      ]}
-    />
+    <>
+      <Button onClick={removeSelected}>Remove Selected</Button>
+      <DataGrid
+        data={data}
+        striped
+        highlightOnHover
+        withGlobalFilter
+        withPagination
+        withColumnFilters
+        withSorting
+        withColumnResizing
+        withRowSelection
+        state={{
+          rowSelection: selected,
+        }}
+        onRowSelectionChange={setSelected}
+        columns={[
+          {
+            accessorKey: 'text',
+            header: 'Text that is too long for a Header',
+            filterFn: stringFilterFn,
+            size: 300,
+            cell: highlightFilterValue,
+          },
+          {
+            header: 'Animal',
+            columns: [
+              { accessorKey: 'cat', filterFn: stringFilterFn },
+              {
+                accessorKey: 'fish',
+                filterFn: stringFilterFn,
+              },
+            ],
+          },
+          {
+            accessorKey: 'city',
+            filterFn: stringFilterFn,
+          },
+          { accessorKey: 'value', filterFn: numberFilterFn },
+          {
+            accessorKey: 'date',
+            cell: (cell) => cell.getValue<Date>()?.toLocaleDateString(),
+            filterFn: dateFilterFn,
+          },
+          {
+            accessorKey: 'bool',
+            filterFn: booleanFilterFn,
+          },
+        ]}
+      />
+    </>
   );
 }
diff --git a/src/DataGrid.tsx b/src/DataGrid.tsx
index b669763..a2cb3a1 100644
--- a/src/DataGrid.tsx
+++ b/src/DataGrid.tsx
@@ -226,7 +226,8 @@ export function DataGrid<TData extends RowData>({
 
   const handleRowSelectionChange: OnChangeFn<RowSelectionState> = useCallback(
     (arg0) => {
-      table.setState((state) => {
+      table.setState(() => {
+        const state = table.getState();
         const next = functionalUpdate(arg0, state.rowSelection);
         onRowSelectionChange && onRowSelectionChange(next);
         return {