mainAE.py

if __name__ == '__main__':
    model = Autoencoder()
    model.trainModel()
    tensor = model.testImage(7777)
    # Reshape
    tensor = torch.reshape(tensor, (28, 28))
    # toImage function
    toImage = torchvision.transforms.ToPILImage()
    # Convert to image
    image = toImage(tensor)
    # Save image
    image.save('After.png')