Pytorch：三维矩阵在第一维上取平均值

import torch
import numpy as np

# ======初始化一个三维矩阵=====
A = torch.ones((6,3,3))

# ======替换三维矩阵里面的值======
A[0] = torch.ones((3,3)) *2
A[1] = torch.ones((3,3)) *5
A[2] = torch.ones((3,3)) *5
A[3] = torch.ones((3,3)) *5
A[4] = torch.ones((3,3)) *5
A[5] = torch.ones((3,3)) *5

# ======要取平均的矩阵索引（第一维当做索引）===========
ids = np.array([2,3,4,5])

# =====得到的结果是一个二维矩阵========
B = torch.mean(A[ids],dim=0)
print(B)

结果：