if len(attacks_requesting_batched_predictions) > 0:
logging.debug("calling native forward with {}".format(len(attacks_requesting_batched_predictions))) // noqa: E501
batched_predictions_args = list(map(np.stack,
zip(*batched_predictions_args)))
batched_predictions_args = list(batched_predictions_args)
// get original shape (//attacks, batch size)
batch_shape = batched_predictions_args[0].shape
After Change
logging.debug("calling native forward with {}".format(len(attacks_requesting_batched_predictions))) // noqa: E501
// we are only interested in the first argument
inputs = [x[0] for x in batched_predictions_args]
// merge individual batches into one larger super-batch
batch_lengths = [len(x) for x in inputs]
batch_splits = np.cumsum(batch_lengths)
inputs = np.concatenate([x for x in inputs])
// split super-batch back into individual batches
