torchvision.utils.make_grid改变图像通道

问题：单通道图像经过torchvision.utils.make_grid后，生成了三通道的图像。

torch.Size([4, 1, 256, 256])  →   torch.Size([3, 518, 1034])

原因：该函数内有一段代码：

 当拼接图像为单通道时，会在通道上将图片叠加三次。

解决办法：将代码修改为：