这里用到的是torchvision中的save_image。
废话不多说，直接来代码：
import torch
from torchvision.utils import save_image
B, C, H, W = 64, 3, 32, 32
input_tensor = torch.randn(B, C, H, W)
save_image(input_tensor, "hh.png", nrow=8)
首先，这里我们要明确
(
B
,
C
,
H
,
W
)
(B, C, H, W)
(B,C,H,W)的意思，我们有64张3通道的图片需要可视化。那么，我设定每排放8张。实现效果如图：
参考链接：https://pytorch.org/vision/stable/generated/torchvision.utils.save_image.html
