content: import torch
  
  
  def generate_grid(n_vox, interval):
      """
      generate grid
      if 3D volume, grid[:,:,x,y,z]  = (x,y,z)
      :param n_vox:
      :param interval:
      :return:
      """
      with torch.no_grad():
          # Create voxel grid
          grid_range = [torch.arange(0, n_vox[axis], interval) for axis in range(3)]
          grid = torch.stack(torch.meshgrid(grid_range[0], grid_range[1], grid_range[2], indexing="ij"))  # 3 dx dy dz
          # ! don't create tensor on gpu; imbalanced gpu memory in ddp mode
          grid = grid.unsqueeze(0).type(torch.float32)  # 1 3 dx dy dz
  
      return grid
  
  
  if __name__ == "__main__":
      import torch.nn.functional as F
      grid = generate_grid([5, 6, 8], 1)
  
      pts = 2 * torch.tensor([1, 2, 3]) / (torch.tensor([5, 6, 8]) - 1) - 1
      pts = pts.view(1, 1, 1, 1, 3)
  
      pts = torch.flip(pts, dims=[-1])
  
      sampled = F.grid_sample(grid, pts, mode='nearest')
  
      print(sampled)
  