Yep, still the same behavior. I ran this cell and it will print out the grads are none . Thanks
import torch 
import torchvision
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import skimage
from zipfile import ZipFile
import h5py 
from imgaug import augmenters as iaa
from pathlib import Path
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from torchvision.utils import make_grid
from torch.utils.data import Dataset, DataLoader, TensorDataset
from torch.utils.data import random_split
from torch.autograd import Variable
from torchsummary import summary

#define the alignnet model
def get_conv(grid_size):
  model = nn.Sequential (
      nn.MaxPool2d (2),
      nn.Conv2d (2, 20, 5),
      nn.ReLU(),
      nn.MaxPool2d (2),
      nn.Conv2d (20, 20, 5),
      nn.ReLU(),
      nn.MaxPool2d (2),
      nn.Conv2d (20, 20, 2),
      nn.ReLU(),
      nn.MaxPool2d (2),
      nn.Conv2d (20, 20, 5),
      nn.ReLU(),
   
  )
  return model
def get_tensor_info(tensor):
  info = []
  for name in ['requires_grad', 'is_leaf', 'retains_grad', 'grad_fn', 'grad']:
    info.append(f'{name}({getattr(tensor, name, None)})')
  return ' '.join(info)
#initialize the differential grid
#the parameter learn offset will define whether or not to learn the offset values during training
def init_grid(grid_size=8):
  #spacing of the grid
  #-1 is because we have a1 = -1 (and thus there are grid_size - 1 "spacing" grids)
  delta = 2/(grid_size-1)
  np_grid = np.arange(grid_size, dtype=float)
  np_grid = np.full_like(np_grid,float(delta))
  ts_grid_x = torch.FloatTensor(np_grid)
  ts_grid_y = torch.FloatTensor(np_grid)
  diff_i_grid_y, diff_i_grid_x = torch.meshgrid(ts_grid_x,ts_grid_y)
  diff_grid = torch.stack([diff_i_grid_x, diff_i_grid_y])
  diff_grid = diff_grid.view(2*grid_size*grid_size)
  return diff_grid

#perform cumsum operation on a 2d batch of inputs
#takes in grid tensors of shape batch x 2 x grid x grid 
#return grid tensors of shape batch x 2 x grid x grid 
def cumsum_2d(grid, grid_offset_x, grid_offset_y):
  batch_size, dim, grid_1, grid_2 = grid.shape
  grid[:,0,:,0] = -1 
  grid[:,1,0,:] = -1 

  Integrated_grid_x = torch.cumsum(grid[:,0], dim = 2) + grid_offset_x
  Integrated_grid_y = torch.cumsum(grid[:,1], dim = 1) + grid_offset_y
  Integrated_grid = torch.stack([Integrated_grid_x, Integrated_grid_y])
  Integrated_grid = Integrated_grid.permute([1,0,2,3])

  return Integrated_grid

class Net(nn.Module):
  def __init__(self, grid_size):
    super().__init__()
    self.conv = get_conv(grid_size)
    self.flatten = nn.Flatten()
    self.linear1 = nn.Sequential(nn.Linear(80,20),nn.ReLU(),)
    self.linear2 = nn.Linear(20, 2*8*8)
    self.upsampler = nn.Upsample(size = [128, 128], mode = 'bilinear')
    self.linear2.bias = nn.Parameter(init_grid(8).view(-1))
    self.linear2.weight.data.fill_(float(0))
    self.grid_offset_x = torch.tensor(float(0), requires_grad=True)
    self.grid_offset_y = torch.tensor(float(0), requires_grad=True)
    self.grid_offset_x = nn.Parameter(self.grid_offset_x)
    self.grid_offset_y = nn.Parameter(self.grid_offset_y)
    self.grid_size = grid_size

#input shape (n could be any numbers):
#x: nx2x128x128
#src_batch: nx128x128
  def forward(self, x, src_batch):
    print(f'X gradient1: {get_tensor_info(x)}')
    x = self.conv(x)
    print(f'X gradient2: {get_tensor_info(x)}')
    x = self.flatten(x)
    print(f'X gradient3: {get_tensor_info(x)}')
    x = self.linear1(x)
    print(f'X gradient4: {get_tensor_info(x)}')
    x = self.linear2(x)
    print(f'X gradient5: {get_tensor_info(x)}')
    #enforce axial monotinicity using the abs operation
    x = torch.abs(x)
    print(f'X gradient after abs(): {get_tensor_info(x)}')
    batch, grid = x.shape
    x = x.view(batch, 2,self.grid_size,self.grid_size)
    #perform the cumsum operation to restore the original grid from the differential grid
    x = cumsum_2d(x, self.grid_offset_x, self.grid_offset_y)
    print(f'X gradient after cumsum(): {get_tensor_info(x)}')
    #Upsample the grid_size x grid_size warp field to image_size x image_size warp field
    x = self.upsampler(x)
    print(f'X gradient after upsampling: {get_tensor_info(x)}')
    x = x.permute(0,2,3,1)

    #calculate target estimation
    x = nn.functional.grid_sample(src_batch.unsqueeze(0).permute([1,0,2,3]), x)

    return x
#target_image: n x128x128
#warped_image: n x128x128 (output of the Net model)
def Total_Loss (target_image, warped_image, grid_size=8, Lambda=1e-5):
  batch,W,H = warped_image.shape
  #print(f'Warp_field gradient: {get_tensor_info(warp_field)}')

  L2_Loss_f = nn.MSELoss()
  L2_Loss = 1/2 * L2_Loss_f(target_image, warped_image)
  
  Total_loss = L2_Loss
  Total_loss.retain_grad()
  print(f'Total_loss gradient: {get_tensor_info(Total_loss)}')

  return Total_loss


#aug_batch, src_batch, tar_batch input form: nx128x128
#run the model for a single epoch
#data_loader can either be a single data_loader or a iterable list of data loaders (if so, source loader should also be a list of data loaders)
def run_epoch(model, optimizer,  aug_batch, src_batch, tar_batch, grid_size=8):

  aug_batch = torch.tensor(aug_batch, dtype=torch.float32)
  tar_batch = torch.tensor(tar_batch, dtype=torch.float32)
  src_batch = torch.tensor(src_batch, dtype=torch.float32)

  #run forward propagation
  tar_est = warp(model, src_batch, aug_batch, grid_size)        
  tar_est = tar_est.squeeze(dim=1)
  total_loss = Total_Loss(tar_batch, tar_est, grid_size, 1e-5)

  total_loss.backward()
  optimizer.step()
  optimizer.zero_grad()




def warp(model, src_batch, aug_batch, grid_size):
  input_image = torch.stack([src_batch, aug_batch])
  input_image  = input_image.permute([1,0,2,3])
  #run the network
  target_est = model.forward(input_image, src_batch)

  return target_est

def run_model(model,src_batch, tar_batch, grid_size=8,):
  optimizer = optim.Adam(model.parameters(), lr=1e-3)
  epoch_loss = []
  for i in range(10):
    run_epoch(model, optimizer, src_batch, tar_batch , grid_size)
   
src_batch = torch.randn(5, 128, 128)
tar_batch = torch.randn(5, 128, 128)

model = Net(8)
run_model(model, src_batch, tar_batch)```