IBA for different input dimensions

[zk-fmf]

Hey, I would like to know how can one adjust the code for different input dimensions.

It is working for (3,W,H) images; I would like to use IBA for inputs of dimension (4,3,W,H).

So the input to my model are 4 different images with 3 color channels.

One obvious thing is to expand batch to input_t.expand(batch_size, -1, -1, -1, -1) instead of (batch_size, -1,-1,-1). What else should be adapted?

Thanks for your help and great interpretability method!