def _construct_unfold_indices(N, C, H, W, kernel_size, stride):
    """
    A utility function to construct indices for torch.unfold (im2col),
    assuming the torch.unfold input `x` to be contiguous
    """

    # Get starting block indices.
    start_idx = np.arange(kernel_size[0])[None, :, None] * W + np.arange(
        kernel_size[1]
    )

    # Generate depth indices.
    channel_index = H * W * np.arange(C)
    start_idx = (channel_index[None, :, None] + np.ravel(start_idx)).reshape(
        (-1, kernel_size[0], kernel_size[1])
    )

    # Get offsetted indices across the height and width of input array.
    row_extent = H - kernel_size[0] + 1
    col_extent = W - kernel_size[1] + 1
    offset_idx = np.arange(0, row_extent, stride[0])[None, :, None] * W + np.arange(0, col_extent, stride[1])
    indices = np.ravel(start_idx)[:, None] + np.ravel(offset_idx)

    # Get batch block indices.
    batch_idx = np.arange(N)[:, None, None] * C * H * W
    indices = batch_idx + indices

    return indices.reshape(-1)

    A utility function to construct indices for torch.unfold (im2col),
    assuming the torch.unfold input `x` to be contiguous
    
coremltools/converters/mil/frontend/torch/ops.py