Consistent chunk size.

climax/main/downscale_infer.py

@@ -75,7 +75,7 @@ if __name__ == '__main__':
             dem = dem.drop_vars(['slope', 'aspect'])
 
         # add dem to set of predictor variables
-        Era5_ds = xr.merge([Era5_ds, dem])
+        Era5_ds = xr.merge([Era5_ds, dem]).chunk(Era5_ds.chunks)
 
     # load pretrained model
     if state_file.exists():

climax/main/downscale_train.py

@@ -101,7 +101,7 @@ if __name__ == '__main__':
 
         # check whether to use slope and aspect
         if not DEM_FEATURES:
-            dem = dem.drop_vars(['slope', 'aspect'])
+            dem = dem.drop_vars(['slope', 'aspect']).chunk(Era5_ds.chunks)
 
         # add dem to set of predictor variables
         Era5_ds = xr.merge([Era5_ds, dem])