How to modify this tensor to feed to make_grid?

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?

My network outputs a 3D tensor with shape (B x C x H x W)

B,C,H,W is a 4D Tensor, that’s what make_grid wants, a batch of 3d images.

Thank smth. I have an unrelated error and confused myself.