diff --git a/flox/xarray.py b/flox/xarray.py
index 11cf706d..10a33e1d 100644
--- a/flox/xarray.py
+++ b/flox/xarray.py
@@ -524,7 +524,7 @@ def wrapper(array, *by, func, skipna, core_dims, **kwargs):
                     obj, xr.Dataset
                 )  # do not re-order dataarrays inside datasets
                 actual[var] = _restore_dim_order(
-                    actual[var], template, by_da[0], no_groupby_reorder=no_groupby_reorder
+                    actual[var].variable, template, by_da[0], no_groupby_reorder=no_groupby_reorder
                 )
 
     if missing_dim: