AssertionError assert I.ndim == 4 and I.shape[1] == 3

Run code ：

import torch
import torchvision
from torch import nn
from torch.nn import Conv2d
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

dataset = torchvision.datasets.CIFAR10("CIFAR10",train=False,transform=torchvision.transforms.ToTensor(),download=True)
#  Be careful dataset in transform The parameter receives an object , So we need to add parentheses , In addition, when using neural network for operation, the data type required is tensor type , therefore transforms Parameters to add .
dataloader = DataLoader(dataset,batch_size=64)

#  Build a simple network 
class Booze(nn.Module):

    #  Inherit nn.Module The initialization 
    def __init__(self):
        super().__init__()
        #  Note that here is to create a global variable, so we need to add a self  When out_channels Far greater than in_channels The original image needs to be expanded , That is to say padding The value of cannot be set to 0 了 , According to the formula 

        self.conv1 = Conv2d(in_channels=3,out_channels=6,kernel_size=(3),stride=1,padding=0)

    #  rewrite forward function 
    def forward(self,x):
        x = self.conv1(x)
        return x


#  Initialize the network 
obj = Booze()
#  Check out the Internet 
print(obj)
''' Booze( (conv1): Conv2d(3, 6, kernel_size=(3, 3), stride=(1, 1)) ) '''

writer = SummaryWriter("logs")
step = 0

for data in dataloader:
    imgs,targets = data
    output = obj(imgs)
    # torch.Size([64, 3, 32, 32]) 64 Zhang 3 passageway 32X32 Pictures of the 
    print(imgs.shape)
    # torch.Size([64, 6, 30, 30]) 64 Zhang 6 passageway 30X30 Pictures of the 
    print(output.shape)

    #  Use tensorboard visualization   Note that multiple images are to be used add_images instead of add_image
    writer.add_images("input",imgs,step)

    #  because output yes 6 The number of channels cannot be displayed , Direct visualization will report an error , So we need to deal with output Conduct reshape reshape When a number is unknown in the second parameter of , You can fill in -1, He will automatically help you calculate , Why is it unknown ？ Because I just don't know how much to fill , fill 64 I'm sure not , Then changing the number of channels is equivalent to cutting out the extra pixels 
    writer.add_images("output",output,step)
    step+=1

writer.close()

The error reported by the running code is as follows ：

Why did it go wrong ？
The reason is that in the neural network we build self.conv1 = Conv2d(in_channels=3,out_channels=6,kernel_size=(3),stride=1,padding=0)
among out_channels=6 The number of channels for outputting pictures is 6 The tunnel ,6 The picture of the number of channels cannot be displayed , Use it directly tensorboard Visual error , The error is in the above code writer.add_images("output",output,step) This line of code , So before executing this line of code, you need to output Conduct reshape,reshape become 3 Picture of the number of channels .
Solution
output = torch.reshape(output,(-1,3,30,30)) This line of code is added to writer.add_images("output",output,step) Before ,reshape When a number is unknown in the second parameter of , You can fill in -1, He will automatically help you calculate , Why is it unknown ？ Because I just don't know how much to fill , fill 64 I'm sure not , Then changing the number of channels is equivalent to cutting out the extra pixels , Put it in batch_size in .

# (-1,3,30,30) = (batch_size,channels,H,W)
output = torch.reshape(output,(-1,3,30,30))
writer.add_images("output",output,step)