def METHOD_NAME(test_case):
    _test_linear_graph_train_with_lr_sch(
        test_case, 21, flow.device("cuda"), _sgd_cosine_fn
    )
    _test_linear_graph_train_with_lr_sch(
        test_case, 21, flow.device("cpu"), _sgd_cosine_fn
    )