diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index db376f2d..60e3cb28 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -40,7 +40,7 @@ jobs:
       - name: Checkout source
         uses: actions/checkout@v4
       - name: Create conda environment
-        uses: mamba-org/provision-with-micromamba@main
+        uses: mamba-org/setup-micromamba@v1
         with:
           cache-downloads: true
           micromamba-version: "latest"
@@ -82,7 +82,7 @@ jobs:
           access_token: ${{ github.token }}
       - uses: actions/checkout@v4
       - name: Create conda environment
-        uses: mamba-org/provision-with-micromamba@v16
+        uses: mamba-org/setup-micromamba@v1
         with:
           cache-downloads: true
           micromamba-version: "latest"
diff --git a/xesmf/frontend.py b/xesmf/frontend.py
index a500112f..2fee239c 100644
--- a/xesmf/frontend.py
+++ b/xesmf/frontend.py
@@ -1044,12 +1044,14 @@ def _init_para_regrid(self, ds_in, ds_out, kwargs):
 
         if self.sequence_out:
             ds_out = ds_out.rename({self.out_horiz_dims[1]: 'x_out'})
-            out_chunks = ds_out.chunks.get('x_out')
         else:
             ds_out = ds_out.rename(
                 {self.out_horiz_dims[0]: 'y_out', self.out_horiz_dims[1]: 'x_out'}
             )
-            out_chunks = [ds_out.chunks.get(k) for k in ['y_out', 'x_out']]
+
+        out_chunks = [ds_out.chunks.get(k) for k in ['y_out', 'x_out']]
+        in_chunks = [ds_in.chunks.get(k) for k in ['y_in', 'x_in']]
+        chunks = out_chunks + in_chunks
 
         # Rename coords to avoid issues in xr.map_blocks
         for coord in list(self.out_coords.keys()):
@@ -1060,7 +1062,7 @@ def _init_para_regrid(self, ds_in, ds_out, kwargs):
         weights_dims = ('y_out', 'x_out', 'y_in', 'x_in')
         templ = sps.zeros((self.shape_out + self.shape_in))
         w_templ = xr.DataArray(templ, dims=weights_dims).chunk(
-            out_chunks
+            chunks
         )  # template has same chunks as ds_out
 
         w = xr.map_blocks(