Return predictions for ``dataset``. Predictions can be generated
        as class labels or probabilities.

        Parameters
        ----------
        dataset : SFrame | SArray | dict
            The audio data to be classified.
            If dataset is an SFrame, it must have a column with the same name as
            the feature used for model training, but does not require a target
            column. Additional columns are ignored.

        output_type : {'probability', 'class', 'probability_vector'}, optional
            Form of the predictions which are one of:

            - 'class': Class prediction. For multi-class classification, this
              returns the class with maximum probability.
            - 'probability': Prediction probability associated with the True
              class (not applicable for multi-class classification)
            - 'probability_vector': Prediction probability associated with each
              class as a vector. Label ordering is dictated by the ``classes``
              member variable.

        verbose : bool, optional
            If True, prints progress updates and model details.

        batch_size : int, optional
            If you are getting memory errors, try decreasing this value. If you
            have a powerful computer, increasing this value may improve performance.

        Returns
        -------
        out : SArray
            An SArray with the predictions.

        See Also
        ----------
        evaluate, classify

        Examples
        ----------
        >>> probability_predictions = model.predict(data, output_type='probability')
        >>> prediction_vector = model.predict(data, output_type='probability_vector')
        >>> class_predictions = model.predict(data, output_type='class')

def predict(self, dataset, output_type='class', verbose=True, batch_size=64):
        """
        Return predictions for ``dataset``. Predictions can be generated
        as class labels or probabilities.

        Parameters
        ----------
        dataset : SFrame | SArray | dict
            The audio data to be classified.
            If dataset is an SFrame, it must have a column with the same name as
            the feature used for model training, but does not require a target
            column. Additional columns are ignored.

        output_type : {'probability', 'class', 'probability_vector'}, optional
            Form of the predictions which are one of:

            - 'class': Class prediction. For multi-class classification, this
              returns the class with maximum probability.
            - 'probability': Prediction probability associated with the True
              class (not applicable for multi-class classification)
            - 'probability_vector': Prediction probability associated with each
              class as a vector. Label ordering is dictated by the ``classes``
              member variable.

        verbose : bool, optional
            If True, prints progress updates and model details.

        batch_size : int, optional
            If you are getting memory errors, try decreasing this value. If you
            have a powerful computer, increasing this value may improve performance.

        Returns
        -------
        out : SArray
            An SArray with the predictions.

        See Also
        ----------
        evaluate, classify

        Examples
        ----------
        >>> probability_predictions = model.predict(data, output_type='probability')
        >>> prediction_vector = model.predict(data, output_type='probability_vector')
        >>> class_predictions = model.predict(data, output_type='class')

        """
        from .._mxnet import _mxnet_utils
        import mxnet as mx

        if not isinstance(dataset, (_tc.SFrame, _tc.SArray, dict)):
            raise TypeError('\'dataset\' parameter must be either an SFrame, SArray or dictionary')

        if isinstance(dataset, dict):
            if(set(dataset.keys()) != {'sample_rate', 'data'}):
                raise ValueError('\'dataset\' parameter is a dictionary but does not appear to be audio data.')
            dataset = _tc.SArray([dataset])
        elif isinstance(dataset, _tc.SFrame):
            dataset = dataset[self.feature]

        if not _is_deep_feature_sarray(dataset) and not _is_audio_data_sarray(dataset):
            raise ValueError('\'dataset\' must be either audio data or audio deep features.')

        if output_type not in ('probability', 'probability_vector', 'class'):
            raise ValueError('\'dataset\' parameter must be either an SFrame, SArray or dictionary')
        if output_type == 'probability' and self.num_classes != 2:
            raise _ToolkitError('Output type \'probability\' is only supported for binary'
                                ' classification. For multi-class classification, use'
                                ' predict_topk() instead.')
        if(batch_size < 1):
            raise ValueError("'batch_size' must be greater than or equal to 1")

        if _is_deep_feature_sarray(dataset):
            deep_features = dataset
        else:
            deep_features = get_deep_features(dataset, verbose=verbose)
        
        deep_features = _tc.SFrame({'deep features': deep_features})
        deep_features = deep_features.add_row_number()
        deep_features = deep_features.stack('deep features', new_column_name='deep features')
        deep_features, missing_ids = deep_features.dropna_split(columns=['deep features'])

        if len(missing_ids) > 0:
            _logging.warning("Unable to make predictions for %d examples because they are less than 975ms in length."
                             % len(missing_ids))

        if batch_size > len(deep_features):
            batch_size = len(deep_features)

        y = []
        for batch in mx.io.NDArrayIter(deep_features['deep features'].to_numpy(), batch_size=batch_size):
            ctx = _mxnet_utils.get_mxnet_context()
            if(len(batch.data[0]) < len(ctx)):
                ctx = ctx[:len(batch.data[0])]

            batch_data = batch.data[0]
            if batch.pad != 0:
                batch_data = batch_data[:-batch.pad]    # prevent batches looping back

            batch_data = mx.gluon.utils.split_and_load(batch_data, ctx_list=ctx, batch_axis=0, even_split=False)

            for x in batch_data:
                forward_output = self._custom_classifier.forward(x)
                y += mx.nd.softmax(forward_output).asnumpy().tolist()
        assert(len(y) == len(deep_features))

        # Combine predictions from multiple frames
        sf = _tc.SFrame({'predictions': y, 'id': deep_features['id']})
        probabilities_sum = sf.groupby('id', {'prob_sum': _tc.aggregate.SUM('predictions')})

        if output_type == 'class':
            predicted_ids = probabilities_sum['prob_sum'].apply(lambda x: _np.argmax(x))
            mappings = self._id_to_class_label
            probabilities_sum['results'] = predicted_ids.apply(lambda x: mappings[x])
        else:
            assert output_type in ('probability', 'probability_vector')
            frame_per_example_count = sf.groupby('id', _tc.aggregate.COUNT())
            probabilities_sum = probabilities_sum.join(frame_per_example_count)
            probabilities_sum['results'] = probabilities_sum.apply(lambda row: [i / row['Count'] for i in row['prob_sum']])

        if len(missing_ids) > 0:
            output_type = probabilities_sum['results'].dtype
            missing_predictions = _tc.SFrame({'id': missing_ids['id'],
                                              'results': _tc.SArray([ None ] * len(missing_ids), dtype=output_type)
                                              })
            probabilities_sum = probabilities_sum[['id', 'results']].append(missing_predictions)

        probabilities_sum = probabilities_sum.sort('id')
        return probabilities_sum['results']