Sure, here is a minimized version of my model. I’ve annotated the input sizes on top of each function definitions. Thanks a lot for your help.


#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).to(DEVICE)
  ts_grid_y = torch.FloatTensor(np_grid).to(DEVICE)
  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