b7e0cfb6d469c3745ae2195908daadea9c64d87e,python/tvm/relay/op/strategy/cuda.py,,roi_align_strategy_cuda,#Any#Any#Any#Any#,944

Before Change

    roi_align cuda strategy
    strategy = _op.OpStrategy()
    layout = attrs.layout
    assert layout == "NCHW", "only support nchw for now"
    strategy.add_implementation(
        wrap_compute_roi_align(topi.vision.rcnn.roi_align_nchw),
        wrap_topi_schedule(topi.cuda.schedule_roi_align),
        name="roi_align_nchw.cuda",

After Change

    strategy = _op.OpStrategy()
    layout = attrs.layout

    if layout == "NCHW":
        strategy.add_implementation(
            wrap_compute_roi_align(topi.vision.rcnn.roi_align_nchw),
            wrap_topi_schedule(topi.cuda.schedule_roi_align),
            name="roi_align_nchw.cuda",
        )
    else:
        assert layout == "NHWC", "layout must be NCHW or NHWC."
        strategy.add_implementation(
            wrap_compute_roi_align(topi.vision.rcnn.roi_align_nhwc),
            wrap_topi_schedule(topi.cuda.schedule_roi_align),
            name="roi_align_nhwc.cuda",
        )
    return strategy


@schedule_roi_pool.register(["cuda", "gpu"])

Instances