tensor[:, :, :2]与tensor[:, :, 2]的区别

结合例子理解

输入

import torch
a = torch.rand(3, 3, 3)
print(a)

b = a[:, :, 2]
print(b)

输出

输入

import torch
a = torch.rand(3, 3, 3)
print(a)

b = a[:, :, :2]
print(b)

输出

总结

[: , :, :number]，假设原来的a维度是（3X3X3），[: , :, :2]表示将原来的维度变为（3X3X2）。

[: , :, number]，假设原来的a维度是（3X3X3），[: , :, 2]表示将只保留标号为2的那一列。