Can I use a pretrained model that returns scores inside a dict rather than directly?

[NishantTharani]

I am trying to create a Classifier with a pretrained model:

classifier = PyTorchClassifier(
    model=my_model,
....
)

The problem with my_model is that if I run out = my_model(x), out is not an array of logits, but it is a dict, and out['logits'] is the array of logits.

Is there any way that I can use such a model with ART, perhaps an option/flag that I've missed in my reading of the documentation? Or do I have to alter the model to return an array of logits directly.

Thanks!

[beat-buesser]

Hi @NishantTharani Thank you very much for using ART!

No, the interfaces are only supporting arrays and tensors. I think a simple solution for you could be to wrap the model into a new model forwarding logits as tensor, e.g.:
(Code example, not tested)

class ModelWrapper(nn.Module):
    def __init__(self, model: torch.nn.Module):
    super().__init__()
    self._model = model

    def forward(self, x):
        result = self._model(x)
        return result["logits"]

and provide an instance of ModelWrapper to ART's tools.

[NishantTharani]

Thank you @beat-buesser , that worked perfectly!