Escrever código hoje encontrou um requisito, ou seja, tenho um array a de tamanho (10, 20, 30, 6) e um índice de índice de tamanho (10, 20, 4), o índice é 4 Um número inteiro não repetitivo de 0 a 29, espero usar esta matriz de idade do índice a para fatiar
a = torch.arange(10 * 20 * 30 * 6).reshape(10, 20, 30, 6)
indice = torch.randint(1, 30, (10, 20, 4))
Usar a[indice] diretamente não pode alcançar o resultado desejado
a[:,:,indice] O tamanho resultante é (10,20,10,20,6), o que não está correto.
Solução
a[torch.arange(a.shape[0])[:, None, None], torch.arange(a.shape[1])[None, :, None], indice]
Ou seja, adicione 2 travessias auxiliares nas duas primeiras dimensões