Batching instead of looping gives incorrect result

I have the following function

def random_sample(feature, pool_idx):
    """
    :param feature: [B, N, d] input features matrix
    :param pool_idx: [B, N', max_num] N' < N, N' is the selected position after pooling
    :return: pool_features = [B, N', d] pooled features matrix
    """
    if len(feature.size()) > 3: # bs*c*p*1
        feature = feature.squeeze(dim=3)  # batch*channel*npoints
    num_neigh = pool_idx.shape[-1]
    d = feature.shape[1]
    batch_size = pool_idx.shape[0]
    pool_idx = pool_idx.reshape(batch_size, -1)  # batch*(npoints,nsamples)
    pool_features = torch.gather(
        feature, 2, pool_idx.unsqueeze(1).repeat(1, feature.shape[1], 1)
    ).contiguous()
    pool_features = pool_features.reshape(batch_size, d, -1, num_neigh)
    pool_features = pool_features.max(dim=3, keepdim=True)[0]  # batch*channel*npoints*1
    return pool_features

I call it with r2p_emb = random_sample(rgb_emb0_tmp,inputs['r2p_ds_nei_idx0'])
Where rgb_emb0_tmp is of size (batch, channel, pts, 1)
and inputs['r2p_ds_nei_idx0'] is of size (batch, sample_pts, neighbors)
and r2p_emb will be of size (batch, channel, sample_pts, 1)

Now my input changes to
rgb_emb0_tmp is of size (batch, views, channel, pts, 1) and
inputs['r2p_ds_nei_idx0'] is of size (batch, views, sample_pts, neighbors)
but I want that each view is processed independently so I can concatenate in the end to get
r2p_emb of shape (batch, channel, views * sample_pts, 1)

r2p_emb = torch.cat([random_sample(rgb_emb0_tmp[:,i,:,:],inputs['r2p_ds_nei_idx0'][:,i,:,:]) for i in range(views)],dim=2)

To vectorize this in my mind I just have to reshape rgb_emb0_tmp to shape (batch * views, channel, pts, 1) and inputs['r2p_ds_nei_idx0'] to shape (batch * views, sample_pts, neighbors) and then finally reshape to shape r2p_emb of shape (batch, channel, views * sample_pts, 1)

This is my thought process. But somehow, somewhere this solutions produces two different end results, I checked it with abs(r2p_emb_loop-r2p_emb_vec).max() which is greater than 0. But it only seems to be the order of the values because the min, mean and max of both tensors is identical.

So where is my solution/my thought process wrong?