La ejecución por lotes en lugar de realizar un bucle produce un resultado incorrecto

Tengo la siguiente función:

def random_sample(feature, pool_idx):
    """
    :param feature: [B, N, d] matriz de características de entrada
    :param pool_idx: [B, N', max_num] N' < N, N' es la posición seleccionada después de realizar pooling
    :return: pool_features = [B, N', d] matriz de características obtenidas por pooling
    """
    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

Lo llamo como r2p_emb = random_sample(rgb_emb0_tmp, inputs['r2p_ds_nei_idx0'])
Donde rgb_emb0_tmp tiene un tamaño de (batch, channel, pts, 1)
y inputs['r2p_ds_nei_idx0'] tiene un tamaño de (batch, sample_pts, neighbors)
y r2p_emb tendrá un tamaño de (batch, channel, sample_pts, 1).

Ahora mi entrada cambia a:
rgb_emb0_tmp tiene un tamaño de (batch, views, channel, pts, 1) y
inputs['r2p_ds_nei_idx0'] tiene un tamaño de (batch, views, sample_pts, neighbors).
Pero quiero que cada vista se procese de manera independiente para poder concatenarlas al final y obtener
r2p_emb de forma (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)

Para vectorizar esto en mi mente, solo tengo que cambiar la forma de rgb_emb0_tmp a (batch * views, channel, pts, 1) y inputs['r2p_ds_nei_idx0'] a (batch * views, sample_pts, neighbors), y luego finalmente volver a dar forma a r2p_emb, que tiene una forma de (batch, channel, views * sample_pts, 1).

Este es mi proceso de pensamiento. Pero de alguna manera, en alguna parte, esta solución produce dos resultados finales diferentes. Lo verifiqué con abs(r2p_emb_loop-r2p_emb_vec).max(), que es mayor que 0. Pero parece ser solo el orden de los valores, porque el mínimo, la media y el máximo de ambos tensores son idénticos.

Entonces, ¿dónde está mal mi solución/mi proceso de pensamiento?