def _try_replace_with_core_upsample(op):
    """
    Inputs:

    op (Operation): op.op_type must be either
    1. `torch_upsample_nearest_neighbor`
    2. `torch_upsample_bilinear`

    Returns:

    True if op can be represented by mb.upsample_nearest_neighbor or mb.upsample_bilinear op in SSA.
    False otherwise
    """
    assert op.op_type in target_ops

    # 2d upsampling
    if op.op_type in ["torch_upsample_nearest_neighbor", "torch_upsample_bilinear"]:
        scales_h = _try_get_upsample_factor(op.output_height.op)
        scales_w = _try_get_upsample_factor(op.output_width.op)

        if scales_h is None or scales_w is None:
            return False

        old_upsample = op.outputs[0]
        block = op.enclosing_block

        if op.op_type == "torch_upsample_nearest_neighbor":
            new_upsample = mb.upsample_nearest_neighbor(
                x=op.x,
                scale_factor_height=scales_h,
                scale_factor_width=scales_w,
                name=op.name,
                before_op=op,
            )
        elif op.op_type == "torch_upsample_bilinear":
            new_upsample = mb.upsample_bilinear(
                x=op.x,
                scale_factor_height=scales_h,
                scale_factor_width=scales_w,
                align_corners=op.align_corners,
                name=op.name,
                before_op=op,
            )
        block.replace_uses_of_var_after_op(anchor_op=op, old_var=old_upsample, new_var=new_upsample)
        block.remove_ops([op])

    return True

    Inputs:

    op (Operation): op.op_type must be either
    1. `torch_upsample_nearest_neighbor`
    2. `torch_upsample_bilinear`

    Returns:

    True if op can be represented by mb.upsample_nearest_neighbor or mb.upsample_bilinear op in SSA.
    False otherwise
    
coremltools/converters/mil/frontend/torch/ssa_passes/torch_upsample_to_core_upsample.py