Complete code. I tried a lot and changes a lot but dimension error remain same
import torch
from torch.utils.data import DataLoader

from model import End_to_end
from torch.autograd import Variable
from loss import loss_function
from grid import save_image
from torchvision import datasets, transforms

CUDA = torch.cuda.is_available()
if CUDA:
    model = End_to_end().cuda()
else:
    model = End_to_end()

EPOCHS = 5

tfs = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
])
testset = datasets.ImageFolder(root="/home/khawar/Desktop/End-to-End_IEEE-TVSCT/test/", transform=tfs)
test_loader = DataLoader(testset, batch_size=8, shuffle=True)
print(test_loader.dataset)
model.load_state_dict(torch.load('./checkpoint/model.pth'))
print(testset.imgs)


def test(epoch):
    model.eval()
    test_loss = 0
    for i, (data, _) in enumerate(test_loader.dataset.imgs):
        data = Variable(data).unsqueeze(0)
        final, residual_img, upscaled_image, com_img, orig_im = model(data.cuda())
        test_loss += loss_function(final, residual_img, upscaled_image, com_img, orig_im).data
        if epoch == EPOCHS and i == 0:
            #             save_image(final.data[0],'reconstruction_final',nrow=8)
            #             save_image(com_img.data[0],'com_img',nrow=8)
            n = min(data.size(0), 6)
            print("saving the image " + str(n))
            comparison = torch.cat([data[:n], final[:n].cpu()])
            comparison = comparison.cpu()
            #             print(comparison.data)
            save_image(com_img[:n].data, 'compressed_' + str(epoch) + '.png', nrow=n)
            save_image(comparison.data, 'reconstruction_' + str(epoch) + '.png', nrow=n)

    test_loss /= len(testset.test_loader)
    print('====> Test set loss: {:.4f}'.format(test_loss))


def save_images():
    epoch = EPOCHS
    model.eval()
    test_loss = 0
    for i, (data, _) in enumerate(test_loader):
        print(data)
        data = Variable(data).unsqueeze(0)
        print(data.shape)
        print(data.size())
        final, residual_img, upscaled_image, com_img, orig_im = model(data.cuda())
        test_loss += loss_function(final, residual_img, upscaled_image, com_img, orig_im).data
        if i == 3:
            #             save_image(final.data[0],'reconstruction_final',nrow=8)
            #             save_image(com_img.data[0],'com_img',nrow=8)
            n = min(data.size(0), 6)
            print("saving the image " + str(n))
            comparison = torch.cat([data[:n], final[:n].cpu()])
            comparison = comparison.cpu()
            #             print(comparison.data)
            save_image(com_img[:1].data, './compressed_image/compressed_' + str(i) + '.png', nrow=n)
            save_image(final[:1].data, './final_image/final_' + str(epoch) + '.png', nrow=n)
            save_image(orig_im[:1].data, './orginal_image/original_' + str(epoch) + '.png', nrow=n)

    test_loss /= len(test_loader.dataset)
    print('====> Test set loss: {:.4f}'.format(test_loss))


save_images()