您可以通过以下方式实现这一点：

groups = []

for i in range(512):
    # Select the i-th embedding from each image
    group = embeddings[:, i, :]
    groups.append(group)

groups = np.array(groups)

另一个优化的解决方案：

groups = np.array([embeddings[:, i, :] for i in range(512)])
groups = np.transpose(groups, (1, 0, 2))