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).

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)