message: Remove additional code related to
diff: import tensorflow as tf
  from absl import flags
  -from autofl.datasets import storage, testing
  +from autofl.datasets import storage
  from autofl.types import FederatedDataset
  FLAGS = flags.FLAGS
  @@ -32,12 +32,6 @@ config = {
  def load_splits(
  dataset_name: str, get_local_datasets_dir=storage.default_get_local_datasets_dir
  ) -> FederatedDataset:
  - federated_dataset = storage.load_splits(
  + return storage.load_splits(
  dataset_name=dataset_name, get_local_datasets_dir=get_local_datasets_dir
  )
  -
  - testing.assert_dataset_origin(
  - keras_dataset=keras_dataset, federated_dataset=federated_dataset
  - )
  -
  - return federated_dataset
  