def format_data(dataset):
    """ data_x, data_y (batch_size, channel, height, width)
    data_x: (1000, 1, 128, 128)
    data_y: (1000, 1, 256, 256)
    """
    data_x, data_y = dataset
    total_batch_size = data_x.shape[0]
    train_index = int(total_batch_size * 0.7)
    valid_index = int(total_batch_size * 0.9)
    np_train_data_x = data_x[:train_index]
    np_valid_data_x = data_x[train_index:valid_index]
    np_test_data_x = data_x[valid_index:]
    np_train_data_y = data_y[:train_index]
    np_valid_data_y = data_y[train_index:valid_index]
    np_test_data_y = data_y[valid_index:]
    return [(np_train_data_x, np_train_data_y), (np_valid_data_x,
        np_valid_data_y), (np_test_data_x, np_test_data_y)]
