And I also have tried to add another con2d to make three channels.But it doesn’t work .
    for data in test_loader:
        # data=data.unsqueeze(0)
        output=model(data)
        x=nn.Conv2d(21,3,1)
        output=x(output)
        print(output.shape[0],output.shape[1],output.shape[2],output.shape[3])
        print(output)
        # output1=output.reshape(42,400,400)
        # print(output1.shape[0],output1.shape[1],output1.shape[2])
        tensor_to_img(output)

The error information is “raise ValueError(‘pic should be 2/3 dimensional. Got {} dimensions.’.format(pic.ndimension()))
ValueError: pic should be 2/3 dimensional. Got 4 dimensions.” as the same last.