torchvision存tensor为图片时报错TypeError: Cannot handle this data type

博主在玩retro上的atari小游戏时，想存一下observation的图片，代码如下：

obs_ = torch.Tensor(obs)
        print(obs_.shape)
        vutils.save_image(obs_, 'kan.png')

张量维度为(250, 160, 3)
结果直接报错
TypeError: Cannot handle this data type (1, 1, 250)
查阅资料后尝试n种解决方法无果，直接深入源码查看：
回溯到倒数第二层：
torch_vision/utils.py/save_img

grid = make_grid(tensor, **kwargs)
    # Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
    ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
    im = Image.fromarray(ndarr)
    im.save(fp, format=format)

报错位置在im = Image.fromarray(ndarr)

经过多次尝试，发现Image.formarray(ndarr)方法，要求ndarr的最后一个维度必须是通道维度。print(ndarr)，发现它的形状是(160, 3, 250).
在中间插入一行，转置改变ndarr的维度位置就行了

	ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()
	ndarr = np.transpose(ndarr, (2, 0, 1))
    im.save(fp, format=format)

转置后的ndarr形状(250, 160, 3)
解决问题。
存出来的图片是灰度图：