I am trying to build a model which can be trained/evaluaed with various combinations of loss functions/metric functions. E.g. I want a model of the form:

```
class Model(nn.Module):
def __init__(self, loss_fn, metric_fn):
...
```

However, since the signature of loss functions (e.g. Cross Entropy) is different than these of a metric (e.g. Accuracy), one should reshape `preds`

and `targets`

according to the loss/metric signature. That is, the following might not always work:

```
loss = loss_fn(input=preds, target=target)
metric = metric_fn(preds=preds, target=target)
```

since the required shapes might be different in `loss_fn`

compared to `metric_fn`

.

This means that the source code should updated each time a different loss and/or metric is to be used. Obviously, this is something we would like to avoid. Is there any solution that I am missing or I need to handle each special case in the source code?