5bce73c2a8b418d5b97e445ba73418adc22802f2,tests/python/relay/test_auto_scheduler_task_extraction.py,,test_task_extraction_cuda,#,71

Before Change

    assert sum(task_weights) == 2

    mod, params = get_network("resnet-18")
    target = tvm.target.Target("cuda")
    tasks, task_weights = auto_scheduler.extract_tasks(mod["main"], params, target)

    assert len(tasks) == 21
    assert sum(task_weights) == 22

After Change

    for layout in ["NCDHW", "NDHWC"]:
        mod, params = get_network("resnet3d-18", layout=layout)
        tasks, task_weights = auto_scheduler.extract_tasks(mod["main"], params, target)

        assert len(tasks) == 21
        assert sum(task_weights) == 22

Instances