Handle model with tensor that has channel first

[AdrienGauffriau]

The explaination does not work for models that ask for the channel as the first dimension of the input tensor (channel, height, width.
When such a model is provided, the output shape of the explanation is (channel x height) instead of (height x width)

[lucashervier]

Hi Adrien!

Would a wrapper such as the one used in this example fix it ?