Hey! I am currently implementing an autoencoder with a simple MNIST dataset. My goals are as follows:

Autoencode the images using the network. 28x28 input and 28x28 output
Test the Model with an input image, and see what the output would be
Finally, I’m looking for it to have a broken image as TRAINING input. I manually break it by covering it with a grey square in the middle. So basically, the network never has seen the MNIST files before corruption, but it should try to hallucinate the image before corruption.

Of course point 3 is very specific, but the first issue I’m currently facing is testing the model itself.
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
from torch.autograd import Variable

def to_img(x):
    x = x.view(x.size(0), 1, 28, 28)
    return x

def plot_sample_img(img, name):
  img = img.view(1, 28, 28)
  save_image(img, './sample_{}.png'.format(name))

class AE(nn.Module):
    def __init__(self, **kwargs):
        super().__init__()
        self.encoder_hidden_layer = nn.Linear(
            in_features=kwargs["input_shape"], out_features=128
        )
        self.encoder_output_layer = nn.Linear(
            in_features=128, out_features=128
        )
        self.decoder_hidden_layer = nn.Linear(
            in_features=128, out_features=128
        )
        self.decoder_output_layer = nn.Linear(
            in_features=128, out_features=kwargs["input_shape"]
        )

    def forward(self, features):
        activation = self.encoder_hidden_layer(features)
        activation = torch.relu(activation)
        code = self.encoder_output_layer(activation)
        code = torch.relu(code)
        activation = self.decoder_hidden_layer(code)
        activation = torch.relu(activation)
        activation = self.decoder_output_layer(activation)
        reconstructed = torch.relu(activation)
        return reconstructed

#  use gpu if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# create a model from `AE` autoencoder class
# load it to the specified device, either gpu or cpu
model = AE(input_shape=784).to(device)

# create an optimizer object
# Adam optimizer with learning rate 1e-3

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# mean-squared error loss
criterion = nn.MSELoss()

def erase_middle(image: torch.Tensor) -> torch.Tensor:
    _, height, width = image.size()
    x_start = width // 2 - 5
    x_end = width // 2 + 5
    y_start = height // 2 - 5
    y_end = height // 2 + 5
    # Using slices achieves the same as the for loops
    image[:, y_start:y_end, x_start:x_end] = 0.5
    return image

transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor(),
                                            torchvision.transforms.Lambda(erase_middle)])

train_dataset = torchvision.datasets.MNIST(
    root="~/torch_datasets", train=True, transform=transform, download=True
)

test_dataset = torchvision.datasets.MNIST(
    root="~/torch_datasets", train=False, transform=transform, download=True
)

train_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=128, shuffle=True, num_workers=4, pin_memory=True
)

test_loader = torch.utils.data.DataLoader(
    test_dataset, batch_size=32, shuffle=False, num_workers=4
)

def imshow(img):
    #img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))



dataiter = iter(train_loader)
images, labels = dataiter.next()
imshow(torchvision.utils.make_grid(images))

epochs = 10
for epoch in range(epochs):
    loss = 0
    for batch_features, _ in train_loader:
        # reshape mini-batch data to [N, 784] matrix
        # load it to the active device
        batch_features = batch_features.view(-1, 784).to(device)
        
        # reset the gradients back to zero
        # PyTorch accumulates gradients on subsequent backward passes
        optimizer.zero_grad()
        
        # compute reconstructions
        outputs = model(batch_features)
        
        # compute training reconstruction loss
        train_loss = criterion(outputs, batch_features)
        
        # compute accumulated gradients
        train_loss.backward()
        
        # perform parameter update based on current gradients
        optimizer.step()
        
        # add the mini-batch training loss to epoch loss
        loss += train_loss.item()

        # saving per iter to allow for resume
        torch.save({
            'epoch': epoch + 1,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss
            }, "/content/model_saved")
    
    # compute the epoch training loss
    loss = loss / len(train_loader)
    
    # display the epoch training loss
    print("epoch : {}/{}, loss = {:.6f}".format(epoch + 1, epochs, loss))

model.eval()
for batch_features,_ in test_loader:
    plt.figure()
    batch_features = batch_features.view(-1, 784).to(device)
    pred = model(Variable(batch_features))


imshow(torchvision.utils.make_grid(pred[0].cpu().detach()[0]))

The last few lines, was my attempt at plotting the images. Could someone guide me how to plot the image from this? Ideally it would show the input and output images next to each other.
And how would I use the grey square in the middle to attempt to “predict” what has been there before? As far as I know, it is really hard to predict something the neural net has never seen, as it could not learn about it before.
Any help would be appreciated!