Fit the visualizer to the the target variables, which must be 1D
        vectors containing discrete (classification) data. Fit has two modes:

        1. Balance mode: if only y_train is specified
        2. Compare mode: if both train and test are specified

        In balance mode, the bar chart is displayed with each class as its own
        color. In compare mode, a side-by-side bar chart is displayed colored
        by train or test respectively.

        Parameters
        ----------
        y_train : array-like
            Array or list of shape (n,) that containes discrete data.

        y_test : array-like, optional
            Array or list of shape (m,) that contains discrete data. If
            specified, the bar chart will be drawn in compare mode.

def fit(self, y_train, y_test=None):
        """
        Fit the visualizer to the the target variables, which must be 1D
        vectors containing discrete (classification) data. Fit has two modes:

        1. Balance mode: if only y_train is specified
        2. Compare mode: if both train and test are specified

        In balance mode, the bar chart is displayed with each class as its own
        color. In compare mode, a side-by-side bar chart is displayed colored
        by train or test respectively.

        Parameters
        ----------
        y_train : array-like
            Array or list of shape (n,) that containes discrete data.

        y_test : array-like, optional
            Array or list of shape (m,) that contains discrete data. If
            specified, the bar chart will be drawn in compare mode.
        """

        # check to make sure that y_train is not a 2D array, e.g. X
        if y_train.ndim == 2:
            raise YellowbrickValueError((
                "fit has changed to only require a 1D array, y "
                "since version 0.9; please see the docs for more info"
            ))

        # Check the target types for the y variables
        self._validate_target(y_train)
        self._validate_target(y_test)

        # Get the unique values from the dataset
        targets = (y_train,) if y_test is None else (y_train, y_test)
        self.classes_ = unique_labels(*targets)

        # Validate the classes with the class names
        if self.labels is not None:
            if len(self.labels) != len(self.classes_):
                raise YellowbrickValueError((
                    "discovered {} classes in the data, does not match "
                    "the {} labels specified."
                ).format(len(self.classes_), len(self.labels)))

        # Determine if we're in compare or balance mode
        self._mode = BALANCE if y_test is None else COMPARE

        # Compute the support values
        if self._mode == BALANCE:
            self.support_ = np.array([
                (y_train == idx).sum() for idx in self.classes_
            ])

        else:
            self.support_ = np.array([
                [
                    (y == idx).sum() for idx in self.classes_
                ]
                for y in targets
            ])

        # Draw the bar chart
        self.draw()

        # Fit returns self
        return self