My network outputs a 3D tensor with shape (B x C x H x W), I understand that torchvision’s make_grid utility takes a 4D tensor.
If I understand right, I want each layer of B(atch) to extend into the 4th dim for make_grid. How can I do this?