def METHOD_NAME(test_case, device):
    input = np.array([[1, 2], [3, 4]])
    index = np.array([[0, 0], [1, 0]])
    np_out = np.take_along_axis(input, index, 1)
    input = flow.tensor(input, dtype=flow.float32, device=flow.device(device))
    index = flow.tensor(index, dtype=flow.int64, device=flow.device(device))
    output = input.gather(1, index)
    test_case.assertTrue(np.array_equal(output.numpy(), np_out))