Изменить размер каналов тензора

Question

У меня есть тензор резака с 3 каналами, и я хочу, чтобы он был 1 каналом (все остальные размеры должны оставаться такими же). Поэтому, если мои текущие размеры torch.Size([6, 3, 512, 512]), я хочу, чтобы оно было torch.Size([6, 1, 512, 512])

Как я могу это сделать?

Answer 1

Решает ли это вашу проблему?

a = torch.ones(6, 3, 512, 512)
b = a[:, 0:1, :, :]
print(b.size()) # torch.Size([6, 1, 5, 5])