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

    generate grid
    if 3D volume, grid[:,:,x,y,z]  = (x,y,z)
    :param n_vox:
    :param interval:
    :return:
    
mv-diffusion/One-2-3-45-by-view/reconstruction/ops/generate_grids.py