# Buggy code
def predict(self, x, batch_size=128, **kwargs):
    """
    Perform prediction for a batch of inputs. Predictions from classifiers should only be aggregated if they all
    have the same type of output (e.g., probabilities). Otherwise, use `raw=True` to get predictions from all
    models without aggregation. The same option should be used for logits output, as logits are not comparable
    between models and should not be aggregated.

    :param x: Test set.
    :type x: `np.ndarray`
    :param raw: Return the individual classifier raw outputs (not aggregated).
    :type raw: `bool`
    :return: Array of predictions of shape `(nb_inputs, nb_classes)`, or of shape
             `(nb_classifiers, nb_inputs, nb_classes)` if `raw=True`.
    :rtype: `np.ndarray`
    """
    if "raw" in kwargs:
        raw = kwargs["raw"]
    else:
        raise ValueError("Missing argument `raw`.")

    preds = np.array(
        [
            self._classifier_weights[i] * self._classifiers[i].predict(x)
            for i in range(self._nb_classifiers)
        ]
    )
    if raw:
        return preds

    # Aggregate predictions only at probabilities level, as logits are not comparable between models
    var_z = np.sum(preds, axis=0)
    return var_z


# Traceback (ValueError)
import torch
from torchvision import models

from art.classifiers import PyTorchClassifier, EnsembleClassifier
from art.attacks import ProjectedGradientDescent, HopSkipJump

# load and preprocess imagenet images
images = load_preprocess_images(...)

resnet18 = models.resnet18(pretrained=True)
alexnet = models.alexnet(pretrained=True)

wrapped_models = []
for model in [resnet18, alexnet]:
loss = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), .1, momentum=0.9, weight_decay=1e-4)
wrapped_model = PyTorchClassifier(model=model, loss=loss, optimizer=optimizer, input_shape=(3,224,224), nb_classes=1000)
wrapped_models.append(wrapped_model)

ensemble = EnsembleClassifier(wrapped_models)
attack = ProjectedGradientDescent(ensemble)
adv_images = attack.generate(x = images)

"""
yields
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-117-766010f0b586> in <module>
1 attack = ProjectedGradientDescent(ensemble, eps=10000)
----> 2 adv_image = attack.generate(x = images)

~/anaconda3/envs/brainscore/lib/python3.6/site-packages/art/attacks/projected_gradient_descent.py in generate(self, x, y, **kwargs)
108
109             # Use model predictions as correct outputs
--> 110             targets = get_labels_np_array(self.classifier.predict(x, batch_size=self.batch_size))
111         else:
112             targets = y

~/anaconda3/envs/brainscore/lib/python3.6/site-packages/art/classifiers/ensemble.py in predict(self, x, batch_size, **kwargs)
114             raw = kwargs['raw']
115         else:
--> 116             raise ValueError('Missing argument `raw`.')
117
118         preds = np.array([self._classifier_weights[i] * self._classifiers[i].predict(x)

ValueError: Missing argument `raw`.

"""

# Fixed code
def predict(self, x, batch_size=128, raw=False, **kwargs):
    """
    Perform prediction for a batch of inputs. Predictions from classifiers should only be aggregated if they all
    have the same type of output (e.g., probabilities). Otherwise, use `raw=True` to get predictions from all
    models without aggregation. The same option should be used for logits output, as logits are not comparable
    between models and should not be aggregated.

    :param x: Test set.
    :type x: `np.ndarray`
    :param raw: Return the individual classifier raw outputs (not aggregated).
    :type raw: `bool`
    :return: Array of predictions of shape `(nb_inputs, nb_classes)`, or of shape
             `(nb_classifiers, nb_inputs, nb_classes)` if `raw=True`.
    :rtype: `np.ndarray`
    """
    preds = np.array(
        [
            self._classifier_weights[i] * self._classifiers[i].predict(x)
            for i in range(self._nb_classifiers)
        ]
    )
    if raw:
        return preds

    # Aggregate predictions only at probabilities level, as logits are not comparable between models
    var_z = np.sum(preds, axis=0)
    return var_z