def METHOD_NAME(img_arr, dim_order='NCHW'):
    """Rearrange image NDArray shape to 'NCHW' or 'NHWC' which
    is valid for MXNet model input.
    Input image NDArray should has dim_order of 'HWC'.
    Parameters
    ----------
    img_arr : NDArray
        Image in NDArray format with shape (channel, width, height)
    dim_order : str
        Output image dimension order. Valid values are 'NCHW' and 'NHWC'
    Returns
    -------
    output : NDArray
        Image in NDArray format with dim_order shape
    """
    assert dim_order in 'NCHW' or dim_order in 'NHWC', "dim_order must be 'NCHW' or 'NHWC'."
    if dim_order == 'NCHW':
        img_arr = mx.nd.transpose(img_arr, (2, 0, 1))
    output = mx.nd.expand_dims(img_arr, axis=0)
    return output