Hi @albanD,
I’m trying to keep the gradient of the tensor in the way below, but the outcome is not as expected. Please kindly check it out if this has any problem that can be solved. Thank you!!!
class S2ConvNet_original(nn.Module):

    def __init__(self):
        super(S2ConvNet_original, self).__init__()

        f1 = 21
        f2 = 40

        self.conv1 = nn.Conv2d(1, f1, 3, stride=3, padding=0)
        self.upsample1 = nn.UpsamplingNearest2d(scale_factor=3)
        self.conv2 = nn.Conv2d(f1, f2, 3, stride=3, padding=0)
        self.upsample2 = nn.UpsamplingNearest2d(scale_factor=3)
        self.fc1 = nn.Linear(in_features=40*60*60, out_features=10)


    def forward(self, x, coor):
        x1 = self.upsample1(x)
        x = preprocess60x60(x, x1, coor)
        x = self.conv1(x)
        x = F.relu(x)
        x2 = self.upsample2(x)
        x = preprocess60x60_21chan(x, x2, coor)
        x = self.conv2(x)
        x = F.relu(x)
        x = self.fc1(x.reshape(x.shape[0], -1))

        return x
"""
This part is for the processing functions
"""
@jit(nopython=False)
def change_patches(patches, image, coor):
    
    a = 0
    for i in range(60): #2160
        for j in range(60): #3840
            coordinate = [i,j]
            indices = np.matrix(coor[(coordinate[0],coordinate[1])])
            patches[a,0,2,1] = image[0,0,indices[1,0],indices[0,0]]
            patches[a,0,2,2] = image[0,0,indices[1,1],indices[0,1]]
            patches[a,0,1,2] = image[0,0,indices[1,2],indices[0,2]]
            patches[a,0,0,2] = image[0,0,indices[1,3],indices[0,3]]
            patches[a,0,0,1] = image[0,0,indices[1,4],indices[0,4]]
            patches[a,0,0,0] = image[0,0,indices[1,5],indices[0,5]]
            patches[a,0,1,0] = image[0,0,indices[1,6],indices[0,6]]
            patches[a,0,2,0] = image[0,0,indices[1,7],indices[0,7]]
            a += 1
    return(patches)

def preprocess60x60(images,newbatch, coor):
    DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    for batch in range(images.shape[0]):
        image = torch.zeros(1,1,60,60)
        image[0,:,:,:] = images[batch,:,:,:]
        
        kc, kh, kw = 1, 3, 3  # kernel size
        dc, dh, dw = 1, 1, 1  # stride
    
        pad = (1,1,1,1)
        paddedimage = F.pad(image, pad).to(DEVICE)        
    
        patches = paddedimage.unfold(1, kc, dc).unfold(2, kh, dh).unfold(3, kw, dw).to(DEVICE)
        unfold_shape = patches.size()
        
        patches = patches.contiguous().view(-1, kc, kh, kw).to(DEVICE)
        numpypatches = patches.cpu().numpy()
        numpyimage = image.cpu().numpy()
        patches = change_patches(numpypatches, numpyimage, coor)
    
        # Reshape back
        patches_orig = torch.from_numpy(patches).view(unfold_shape)
        output_c = unfold_shape[1] * unfold_shape[4]
        output_h = unfold_shape[2] * unfold_shape[5]
        output_w = unfold_shape[3] * unfold_shape[6]
        patches_orig = patches_orig.permute(0, 1, 4, 2, 5, 3, 6).contiguous()
        patches_orig = patches_orig.view(1, output_c, output_h, output_w)
        newbatch[batch,:,:,:] = patches_orig
    return newbatch

def preprocess60x60_21chan(images, newbatch, coor):
    DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    newimage = torch.zeros(1,images.shape[1],180,180).to(DEVICE)

    for batch in range(images.shape[0]):
        image = torch.zeros(1,images.shape[1],60,60)
        image[0,:,:,:] = images[batch,:,:,:]
        for channel in range(21):
            
            imagei = torch.zeros(1,1,60,60)
            imagei[0,0,:,:] = image[0,channel,:,:]

            kc, kh, kw = 1, 3, 3  # kernel size
            dc, dh, dw = 1, 1, 1  # stride
        
            pad = (1,1,1,1)
            paddedimage = F.pad(imagei, pad).to(DEVICE)        
        
            patches = paddedimage.unfold(1, kc, dc).unfold(2, kh, dh).unfold(3, kw, dw).to(DEVICE)
            unfold_shape = patches.size()
            
            patches = patches.contiguous().view(-1, kc, kh, kw)
            
            numpypatches = patches.detach().cpu().numpy()
            numpyimage = imagei.detach().cpu().numpy()
            patches = change_patches(numpypatches, numpyimage, coor)
    
            # Reshape back
            patches_orig = torch.from_numpy(patches).view(unfold_shape)
            output_c = unfold_shape[1] * unfold_shape[4]
            output_h = unfold_shape[2] * unfold_shape[5]
            output_w = unfold_shape[3] * unfold_shape[6]
            patches_orig = patches_orig.permute(0, 1, 4, 2, 5, 3, 6).contiguous()
            patches_orig = patches_orig.view(1, output_c, output_h, output_w)
            newimage[0,channel,:,:] = patches_orig
        
        newbatch[batch,:,:,:] = newimage
    return newbatch

"""
The main part for common training and testing
"""
def main(network):
    train_loader, test_loader, train_dataset, _ = load_data( MNIST_PATH, BATCH_SIZE)
    with open('dict60x60.pkl', 'rb') as fp:
        coor = pickle.load(fp)
    if network == 'original':
        classifier = S2ConvNet_original()
    elif network == 'deep':
        classifier = S2ConvNet_deep()
    else:
        raise ValueError('Unknown network architecture')
    classifier.to(DEVICE)

    print("#params", sum(x.numel() for x in classifier.parameters()))

    criterion = nn.CrossEntropyLoss()
    criterion = criterion.to(DEVICE)

    optimizer = torch.optim.Adam(
        classifier.parameters(),
        lr=LEARNING_RATE)


    for epoch in range(NUM_EPOCHS):
        for i, (images, labels) in enumerate(train_loader):
            classifier.train()

            images = images.to(DEVICE)
            #print(images.size())
            
            labels = labels.to(DEVICE)
        
            optimizer.zero_grad()
            outputs = classifier(images, coor)
            loss = criterion(outputs, labels)
            loss.backward()

            optimizer.step()

            print('\rEpoch [{0}/{1}], Iter [{2}/{3}] Loss: {4:.4f}'.format(
                epoch+1, NUM_EPOCHS, i+1, len(train_dataset)//BATCH_SIZE,
                loss.item()), end="")
        print("")
        correct = 0
        total = 0
        for images, labels in test_loader:

            classifier.eval()

            with torch.no_grad():
                images = images.to(DEVICE)
                labels = labels.to(DEVICE)

                outputs = classifier(images, coor)
                _, predicted = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (predicted == labels).long().sum().item()

        print('Test Accuracy: {0}'.format(100 * correct / total))

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--network",
                        help="network architecture to use",
                        default='original',
                        choices=['original', 'deep'])
    args = parser.parse_args()

    main(args.network)