# Buggy code
def fit_before_convert(self, dataset):
    # If in tf.data.Dataset, must be encoded already.
    if isinstance(dataset, tf.data.Dataset):
        if not self.num_classes:
            shape = dataset.take(1).shape[1]
            if shape == 1:
                self.num_classes = 2
            else:
                self.num_classes = shape
        return
    if isinstance(dataset, pd.DataFrame):
        dataset = dataset.values
    if isinstance(dataset, pd.Series):
        dataset = dataset.values.reshape(-1, 1)
    # Not label.
    if len(dataset.flatten()) != len(dataset):
        self.num_classes = dataset.shape[1]
        return
    labels = set(dataset.flatten())
    if self.num_classes is None:
        self.num_classes = len(labels)
    if self.num_classes == 2:
        self.label_encoder = encoders.LabelEncoder()
    elif self.num_classes > 2:
        self.label_encoder = encoders.OneHotEncoder()
    elif self.num_classes < 2:
        raise ValueError(
            "Expect the target data for {name} to have "
            "at least 2 classes, but got {num_classes}.".format(
                name=self.name, num_classes=self.num_classes
            )
        )
    self.label_encoder.fit(dataset)


# Traceback (AttributeError)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-4-a90e48684d55> in <module>
----> 1 ak0.fit(mnist_train, epochs=10)

~/anaconda3/lib/python3.7/site-packages/autokeras/tasks/image.py in fit(self, x, y, epochs, callbacks, validation_split, validation_data, **kwargs)
119                     validation_split=validation_split,
120                     validation_data=validation_data,
--> 121                     **kwargs)
122
123

~/anaconda3/lib/python3.7/site-packages/autokeras/auto_model.py in fit(self, x, y, batch_size, epochs, callbacks, validation_split, validation_data, **kwargs)
229             y=y,
230             validation_data=validation_data,
--> 231             validation_split=validation_split)
232
233         # Process the args.

~/anaconda3/lib/python3.7/site-packages/autokeras/auto_model.py in _prepare_data(self, x, y, validation_data, validation_split)
303         # TODO: Handle other types of input, zip dataset, tensor, dict.
304         # Prepare the dataset.
--> 305         dataset = self._process_xy(x, y, True)
306         if validation_data:
307             self._split_dataset = False

~/anaconda3/lib/python3.7/site-packages/autokeras/auto_model.py in _process_xy(self, x, y, fit)
291
292         x = self._process_x(x, fit)
--> 293         y = self._process_y(y, fit)
294
295         return tf.data.Dataset.zip((x, y))

~/anaconda3/lib/python3.7/site-packages/autokeras/auto_model.py in _process_y(self, y, fit)
267         for data, head, adapter in zip(y, self._heads, self._output_adapters):
268             if fit:
--> 269                 data = adapter.fit_transform(data)
270             else:
271                 data = adapter.transform(data)

~/anaconda3/lib/python3.7/site-packages/autokeras/engine/adapter.py in fit_transform(self, dataset)
66     def fit_transform(self, dataset):
67         self.check(dataset)
---> 68         self.fit_before_convert(dataset)
69         dataset = self.convert_to_dataset(dataset)
70         self.fit(dataset)

~/anaconda3/lib/python3.7/site-packages/autokeras/adapters/output_adapter.py in fit_before_convert(self, dataset)
65         if isinstance(dataset, tf.data.Dataset):
66             if not self.num_classes:
---> 67                 shape = dataset.take(1).shape[1]
68                 if shape == 1:
69                     self.num_classes = 2

AttributeError: 'TakeDataset' object has no attribute 'shape'

# Fixed code
def fit_before_convert(self, dataset):
    # If in tf.data.Dataset, must be encoded already.
    if isinstance(dataset, tf.data.Dataset):
        if not self.num_classes:
            shape = utils.dataset_shape(dataset)[0]
            # Single column with 0s and 1s.
            if shape == 1:
                self.num_classes = 2
            else:
                self.num_classes = shape
        return
    if isinstance(dataset, pd.DataFrame):
        dataset = dataset.values
    if isinstance(dataset, pd.Series):
        dataset = dataset.values.reshape(-1, 1)
    # Not label.
    if len(dataset.flatten()) != len(dataset):
        self.num_classes = dataset.shape[1]
        return
    labels = set(dataset.flatten())
    if self.num_classes is None:
        self.num_classes = len(labels)
    if self.num_classes == 2:
        self.label_encoder = encoders.LabelEncoder()
    elif self.num_classes > 2:
        self.label_encoder = encoders.OneHotEncoder()
    elif self.num_classes < 2:
        raise ValueError(
            "Expect the target data for {name} to have "
            "at least 2 classes, but got {num_classes}.".format(
                name=self.name, num_classes=self.num_classes
            )
        )
    self.label_encoder.fit(dataset)