diff --git a/src/rpa_gw_im_time_util.F b/src/rpa_gw_im_time_util.F
index 6243bed9f6..605757dcbd 100644
--- a/src/rpa_gw_im_time_util.F
+++ b/src/rpa_gw_im_time_util.F
@@ -110,8 +110,8 @@ CONTAINS
       INTEGER :: handle, icol_global, imo, irow_global, min_bsize, min_bsize_mo, nkind, &
          nmo_blk_gw, npcols, nprows, size_MO, unit_nr_prv
       INTEGER, ALLOCATABLE, DIMENSION(:)                 :: dist1, dist2, dist3, sizes_AO, &
-                                                            sizes_AO_split, sizes_MO, sizes_RI, &
-                                                            sizes_RI_split, tmp
+                                                            sizes_AO_split, sizes_MO, sizes_MO_1, &
+                                                            sizes_RI, sizes_RI_split, tmp
       INTEGER, DIMENSION(2)                              :: pdims_2d
       INTEGER, DIMENSION(2, 1)                           :: bounds
       INTEGER, DIMENSION(3)                              :: pdims
@@ -122,7 +122,7 @@ CONTAINS
       TYPE(atomic_kind_type), DIMENSION(:), POINTER      :: atomic_kind_set
       TYPE(cp_fm_type), POINTER                          :: fm_mat_mo_coeff_gw
       TYPE(dbcsr_distribution_type)                      :: dist, dist_templ
-      TYPE(dbcsr_t_pgrid_type)                           :: pgrid_2d, pgrid_ic, pgrid_MO
+      TYPE(dbcsr_t_pgrid_type)                           :: pgrid_2d, pgrid_AO, pgrid_ic, pgrid_MO
       TYPE(dbcsr_t_type) :: mo_coeff_gw_t, mo_coeff_gw_t_tmp, t_3c_overl_int_ao_ao, &
          t_3c_overl_int_ao_mo, t_3c_overl_int_mo_ao, t_3c_overl_int_mo_mo
       TYPE(dbcsr_type)                                   :: mat_mo_coeff_gw_reflected_norm, &
@@ -177,8 +177,9 @@ CONTAINS
       min_bsize = qs_env%mp2_env%ri_rpa_im_time%min_bsize
       min_bsize_mo = qs_env%mp2_env%ri_rpa_im_time%min_bsize_mo
 
-      ALLOCATE (sizes_MO(nmo))
       CALL split_block_sizes([gw_corr_lev_occ + gw_corr_lev_virt], sizes_MO, min_bsize_mo)
+      ALLOCATE (sizes_MO_1(nmo))
+      sizes_MO_1(:) = 1
 
       nmo_blk_gw = SIZE(sizes_MO)
       CALL move_alloc(sizes_MO, tmp)
@@ -198,6 +199,10 @@ CONTAINS
 
       DEALLOCATE (basis_set_ao, basis_set_ri_aux)
 
+      pdims = 0
+      CALL dbcsr_t_pgrid_create(para_env%group, pdims, pgrid_AO, &
+                                tensor_dims=[SIZE(sizes_RI_split), SIZE(sizes_AO_split), SIZE(sizes_AO_split)])
+
       pdims_2d = 0
       CALL mp_dims_create(para_env%num_pe, pdims_2d)
 
@@ -208,14 +213,14 @@ CONTAINS
 
       pdims_2d = 0
       CALL dbcsr_t_pgrid_create(para_env%group, pdims_2d, pgrid_2d, &
-                                tensor_dims=[SIZE(sizes_AO_split), 1])
+                                tensor_dims=[SIZE(sizes_AO_split), nmo])
 
-      CALL create_3c_tensor(t_3c_overl_int_ao_ao, dist1, dist2, dist3, pgrid_MO, &
+      CALL create_3c_tensor(t_3c_overl_int_ao_ao, dist1, dist2, dist3, pgrid_AO, &
                             sizes_RI_split, sizes_AO_split, sizes_AO_split, [1, 2], [3], name="(RI AO | AO)")
       DEALLOCATE (dist1, dist2, dist3)
       CALL dbcsr_t_copy(t_3c_overl_int(1, 1), t_3c_overl_int_ao_ao)
-      CALL create_3c_tensor(t_3c_overl_int_ao_mo, dist1, dist2, dist3, pgrid_MO, &
-                            sizes_RI_split, sizes_AO_split, sizes_MO, [1, 2], [3], name="(RI AO | MO)")
+      CALL create_3c_tensor(t_3c_overl_int_ao_mo, dist1, dist2, dist3, pgrid_AO, &
+                            sizes_RI_split, sizes_AO_split, sizes_MO_1, [1, 2], [3], name="(RI AO | MO)")
       DEALLOCATE (dist1, dist2, dist3)
 
       CALL create_3c_tensor(t_3c_overl_int_gw_RI, dist1, dist2, dist3, pgrid_MO, &
@@ -226,9 +231,10 @@ CONTAINS
                             sizes_AO_split, sizes_RI_split, sizes_MO, [1], [2, 3], name="(AO | RI MO)")
       DEALLOCATE (dist1, dist2, dist3)
 
+      CALL dbcsr_t_pgrid_destroy(pgrid_AO)
       CALL dbcsr_t_pgrid_destroy(pgrid_MO)
 
-      CALL create_2c_tensor(mo_coeff_gw_t, dist1, dist2, pgrid_2d, sizes_AO_split, sizes_MO, name="(AO|MO)")
+      CALL create_2c_tensor(mo_coeff_gw_t, dist1, dist2, pgrid_2d, sizes_AO_split, sizes_MO_1, name="(AO|MO)")
       DEALLOCATE (dist1, dist2)
       CALL dbcsr_t_pgrid_destroy(pgrid_2d)
 
@@ -253,17 +259,17 @@ CONTAINS
       IF (do_ic_model) THEN
          pdims = 0
          CALL dbcsr_t_pgrid_create(para_env%group, pdims, pgrid_ic, &
-                                   tensor_dims=[SIZE(sizes_RI_split), nmo_blk_gw, nmo_blk_gw])
+                                   tensor_dims=[SIZE(sizes_RI_split), nmo, nmo])
 
          CALL create_3c_tensor(t_3c_overl_int_mo_ao, dist1, dist2, dist3, pgrid_ic, &
-                               sizes_RI_split, sizes_MO, sizes_AO_split, [1, 2], [3], name="(RI MO | AO)")
+                               sizes_RI_split, sizes_MO_1, sizes_AO_split, [1, 2], [3], name="(RI MO | AO)")
          DEALLOCATE (dist1, dist2, dist3)
          CALL create_3c_tensor(t_3c_overl_int_mo_mo, dist1, dist2, dist3, pgrid_ic, &
-                               sizes_RI_split, sizes_MO, sizes_MO, [1, 2], [3], name="(RI MO | MO)")
+                               sizes_RI_split, sizes_MO_1, sizes_MO_1, [1, 2], [3], name="(RI MO | MO)")
          DEALLOCATE (dist1, dist2, dist3)
          CALL dbcsr_t_create(t_3c_overl_int_mo_mo, t_3c_overl_nnP_ic)
          CALL create_3c_tensor(t_3c_overl_nnP_ic_reflected, dist1, dist2, dist3, pgrid_ic, &
-                               sizes_RI_split, sizes_MO, sizes_MO, [1], [2, 3], name="(RI | MO MO)")
+                               sizes_RI_split, sizes_MO_1, sizes_MO_1, [1], [2, 3], name="(RI | MO MO)")
          DEALLOCATE (dist1, dist2, dist3)
 
          CALL dbcsr_t_pgrid_destroy(pgrid_ic)