How can I get the same image from the input?

My Code is listed below:

def train(epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(trainloader):
        #print str(target.numpy())
        h = data.shape[2]
        w = data.shape[3]
        c = 3
        output = data.squeeze(0)
        output = output.numpy()
        output = np.resize(output,(h,w,c))
        plt.imshow(output)
        plt.show()
        exit(0)

Actually, I want to get a output whose locations of each pixels are as same as my original input image, and however my final output is very different from my input image. So, how can i solve this problem?

Thank you for your help.

Maybe you are looking for np.transpose(output, [1, 2 ,0])… But It seems that you are not fwding through the model, so I’m not sure what you want to do with output.