docs: def search(
          self,
          search_space,
          valid_data,
          init_args=[],
          train_args=[],
          init_kwargs={},
          train_kwargs={},
          module_args={},
          module_kwargs={},
          max_search=None,
          shuffle=True,
          verbose=True,
          seed=None,
          **score_kwargs,
      ):
          """
          Performs hyperband search according to the generated schedule.
  
          At the beginning of each bracket, we generate a
          list of random configurations and perform
          successive halving on it; we repeat this process
          for the number of brackets in the schedule.
  
          Args:
              init_args: (list) positional args for initializing the model
              train_args: (list) positional args for training the model
              valid_data: a tuple of Tensors (X,Y), a Dataset, or a DataLoader of
                  X (data) and Y (labels) for the dev split
              search_space: see ModelTuner's config_generator() documentation
              max_search: see ModelTuner's config_generator() documentation
              shuffle: see ModelTuner's config_generator() documentation
  
          Returns:
              best_model: the highest performing trained model found by Hyperband
              best_config: (dict) the config corresponding to the best model
  
          Note: Initialization is performed by ModelTuner instead of passing a
          pre-initialized model so that tuning may be performed over all model
          parameters, including the network architecture (which is defined before
          the train loop).
          """
          self._clear_state(seed)
          self.search_space = search_space
  
          # Loop over each bracket
          n_models_scored = 0
          for bracket_index, bracket in enumerate(self.hyperband_schedule):
  
              # Sample random configurations to seed SuccessiveHalving
              n_starting_configurations, _ = bracket[0]
              configurations = list(
                  self.config_generator(
                      search_space,
                      max_search=n_starting_configurations,
                      rng=self.rng,
                      shuffle=True,
                  )
              )
  
              # Successive Halving
              for band_index, (n_i, r_i) in enumerate(bracket):
  
                  assert len(configurations) <= n_i
  
                  # Evaluate each configuration for r_i epochs
                  scored_configurations = []
                  for i, configuration in enumerate(configurations):
  
                      cur_model_index = n_models_scored
  
                      # Set epochs of the configuration
                      configuration["n_epochs"] = r_i
  
                      # Train model and get the score
                      score, model = self._test_model_config(
                          f"{band_index}_{i}",
                          configuration,
                          valid_data,
                          init_args=init_args,
                          train_args=train_args,
                          init_kwargs=init_kwargs,
                          train_kwargs=train_kwargs,
                          module_args=module_args,
                          module_kwargs=module_kwargs,
                          verbose=verbose,
                          **score_kwargs,
                      )
  
                      # Add score and model to list
                      scored_configurations.append(
                          (score, cur_model_index, configuration)
                      )
                      n_models_scored += 1
  
                  # Sort scored configurations by score
                  scored_configurations.sort(key=lambda x: x[0], reverse=True)
  
                  # Successively halve the configurations
                  if band_index + 1 < len(bracket):
                      n_to_keep, _ = bracket[band_index + 1]
                      configurations = [x[2] for x in scored_configurations][:n_to_keep]
  
          print("=" * 60)
          print(f"[SUMMARY]")
          print(f"Best model: [{self.best_index}]")
          print(f"Best config: {self.best_config}")
          print(f"Best score: {self.best_score}")
          print("=" * 60)
  
          # Return best model
          return self._load_best_model(clean_up=True)