def group_by_label(datapoints, num_labels):
  """Groups a list of datapoints by the classification label.

  Args:
    datapoints: (list of dicts) list of datapoints, each being a dict containing
      keys 'image' and 'label'.
    num_labels: (int) number of labels.

  Returns:
    A list of lists, containing all elements of `datapoints` grouped by
    the 'label' key.
  """
  data_grouped_by_label = [[] for _ in range(num_labels)]

  for datapoint in datapoints:
    data_grouped_by_label[datapoint['label']].append(datapoint)

  return data_grouped_by_label