You can slice specific parts then concatenate them to one tensor
>>> a = torch.rand((1,3,3,5))
tensor([[[[0.1562, 0.1055, 0.5032, 0.0966, 0.2088],
          [0.8520, 0.3314, 0.6602, 0.4182, 0.7371],
          [0.5483, 0.2338, 0.9695, 0.6638, 0.9435]],

         [[0.2520, 0.9377, 0.7031, 0.9447, 0.9292],
          [0.3504, 0.1614, 0.0280, 0.7294, 0.4252],
          [0.6095, 0.9667, 0.6628, 0.9642, 0.2766]],

         [[0.8955, 0.5914, 0.3912, 0.8431, 0.9219],
          [0.4871, 0.5351, 0.9024, 0.3061, 0.1959],
          [0.2931, 0.2810, 0.7195, 0.8711, 0.6927]]]])
>>> c = a[:,:,0,:]
tensor([[[0.1562, 0.1055, 0.5032, 0.0966, 0.2088],
         [0.2520, 0.9377, 0.7031, 0.9447, 0.9292],
         [0.8955, 0.5914, 0.3912, 0.8431, 0.9219]]])
>>> d = a[:,:,1,:3]
tensor([[[0.8520, 0.3314, 0.6602],
         [0.3504, 0.1614, 0.0280],
         [0.4871, 0.5351, 0.9024]]])
>>> e = a[:,:,2,:2]
tensor([[[0.5483, 0.2338],
         [0.6095, 0.9667],
         [0.2931, 0.2810]]])
>>> f = torch.cat((c,d,e),2)
tensor([[[0.1562, 0.1055, 0.5032, 0.0966, 0.2088, 0.8520, 0.3314, 0.6602,
          0.5483, 0.2338],
         [0.2520, 0.9377, 0.7031, 0.9447, 0.9292, 0.3504, 0.1614, 0.0280,
          0.6095, 0.9667],
         [0.8955, 0.5914, 0.3912, 0.8431, 0.9219, 0.4871, 0.5351, 0.9024,
          0.2931, 0.2810]]])

I think slices and torch.cat should keep the gradients