Hello,
I want to change end of encoder layer after training phase and in test phase, as any column and row in feature map be given shifts. this is my code but I dont know how to apply these changes. please help me.
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import matplotlib.pyplot as plt
from torch.utils.data.sampler import SubsetRandomSampler
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

#%%

#Converting data to torch.FloatTensor
transform = transforms.ToTensor()


# Download the training and test datasets
train_data = datasets.MNIST(root='data', train=True, download=True, transform=transform)

test_data = datasets.MNIST(root='data', train=False, download=True, transform=transform)

train_loader = torch.utils.data.DataLoader(train_data, batch_size=32, num_workers=0)
test_loader = torch.utils.data.DataLoader(test_data, batch_size=32, num_workers=0)

#train_loader=torch.to

#%%

#Define the Convolutional Autoencoder
class ConvAutoencoder(nn.Module):
   def __init__(self):
       super(ConvAutoencoder, self).__init__()
      
       #Encoder
       self.conv1 = nn.Conv2d(1, 16, 3, stride=2, padding=1)
       self.conv2 = nn.Conv2d(16, 8, 3, stride=2, padding=1)
       self.conv3 = nn.Conv2d(8,8,3)
   
       #Decoder
       self.conv4 = nn.ConvTranspose2d(8, 8, 3)
       self.conv5 = nn.ConvTranspose2d(8, 16, 3, stride=2, padding=1, output_padding=1)
       self.conv6 = nn.ConvTranspose2d(16, 1, 3, stride=2, padding=1, output_padding=1)

   def forward(self, x):
       x = F.relu(self.conv1(x))      
       x = F.relu(self.conv2(x))
       x = F.relu(self.conv3(x))  
       #mask = torch.cat([torch.ones([8,8,2,3]), torch.zeros([8,8,1,3])], 2)
       #x=torch.multiply(x,mask)
      # x[7,7,:,:]=x[6,6,:,:]
       x = F.relu(self.conv4(x))
       x = F.relu(self.conv5(x))
       x = F.relu(self.conv6(x))

       return x

#Instantiate the model
model = ConvAutoencoder()
print(model)

#%%
def train(model, num_epochs=20, batch_size=64, learning_rate=1e-3):
   torch.manual_seed(42)
   criterion = nn.MSELoss() # mean square error loss
   optimizer = torch.optim.Adam(model.parameters(),
                                lr=learning_rate, 
                                weight_decay=1e-5) # <--
  # train_loader =train_loader;

   outputs = []
   for epoch in range(num_epochs):
       for data in train_loader:
           img, _ = data
           recon = model(img)
           loss = criterion(recon, img)
           loss.backward()
           optimizer.step()
           optimizer.zero_grad()

       print('Epoch:{}, Loss:{:.4f}'.format(epoch+1, float(loss)))
       outputs.append((epoch, img, recon),)
   return outputs
#%%

#test_image = test_loader.open(test_image_name).convert('RGB')

model =  ConvAutoencoder()
max_epochs =20
outputs = train(model, num_epochs=max_epochs)
#%%

for k in range(0, max_epochs, 9):
   plt.figure(figsize=(9, 2))
   imgs = outputs[k][1].detach().numpy()
   recon = outputs[k][2].detach().numpy()
   for i, item in enumerate(imgs):
       if i >= 9: break
       plt.subplot(2, 9, i+1)
       plt.imshow(item[0])
       
   for i, item in enumerate(recon):
       if i >= 9: break
       plt.subplot(2, 9, 9+i+1)
       plt.imshow(item[0])
       
       #%%
  
       
a=(ConvAutoencoder().conv3.weight)
print("conv1 filters: ",a.data.size())


#%%
#test phase        

with torch.no_grad():                                                                     #3
   for data in test_loader:    
        data = torch.tensor(data[0])    
        a=(ConvAutoencoder().conv3.weight)
        a0=a[0,0,:,:]
        a1=a[1,1,:,:]
        a2=a[2,2,:,:]
        a3=a[3,3,:,:]
        a4=a[4,4,:,:]
        a5=a[5,5,:,:]
        a6=a[6,6,:,:]
        a7=a[7,7,:,:]
        a0=a1
        a1=a2
        a2=a3
        a3=a4
        a4=a5
        a5=a6
        a6=a7
        a7=a0
        a=[a0,
           a1,
           a2,
           a3,
           a4,
           a5,
           a6,
           a7]
        output = model(data)     

   plt.imshow(output[0,0])