Pytorch tensor.sum()

Pytorch 中 tensor.sum(axis) 会按照指定的维度进行求和，但是对于 N * C * H * W 这种类型的输入，如果通过 sum 获得 N * C * 1 * 1 的输出该怎么做？

a = torch.ones(size=(1, 3, 4, 4))  # 全1矩阵
b = a.sum(axis=[2, 3], keepdim=False)
print(b.shape)    # (1, 3)
print(b)          # [[16, 16, 16]] 

a = torch.ones(size=(1, 3, 4, 4))  # 全1矩阵
b = a.mean(axis=[2, 3], keepdim=True)
print(b.shape)    # (1, 3, 1, 1)
print(b)          # [[ [[1]], [[1]], [[1]] ]]

注意: axis参数可以同时接收多个维度，并可同时在这些维度上进行指定操作。