Repository: pytorch/ignite

Issue:
`model_transform` has been added to `create_supervised_trainer` and `create_supervised_evaluator` but its not used by `create_supervised_evaluator`
## 🐛 Bug description

https://github.com/pytorch/ignite/pull/2848 added `model_transform` to both `create_supervised_trainer` and `create_supervised_evaluator` but its not used by `create_supervised_evaluator`

https://github.com/pytorch/ignite/blob/65ac304cf1e4cd5004a6ae33dfe075beb9ce8a48/ignite/engine/__init__.py#L692

For consistency it should be implemented in both (i.e. `supervised_evaluation_step` and `supervised_evaluation_step_amp`)- otherwise, evaluation fails if the model returns a dictionary for example.





Discussion:
Thanks a lot for reporting @david-waterworth !
Technically, `model_transform` argument shouldn't be in `create_supervised_evaluator` as it is redundant with `output_transform` argument and I oversaw it in the last reviewing round of the PR #2848 (cc @guptaaryan16).

@david-waterworth do you think we still need it or `output_transform` is enough ?

```python
    def evaluate_step(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
        model.eval()
        with torch.no_grad():
            x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
            output = model(x)
            y_pred = model_transform(output)   # <--- may seem a bit redundant
            return output_transform(x, y, y_pred)
```

@vfdev-5 `output_transform` is sufficient. Having both has some benefits though as you can define a single model transform and use it for both the trainer and evaluator so its not totally redundent.

```python
    # model output is tuple, need to extract first element
    model_transform = lambda output: output[0]

    trainer = create_supervised_trainer(
        model,
        optimizer,
        model_transform = model_transform,
        loss_fn=loss_fn,
        prepare_batch=prepare_batch, 
        device=params['train']['device'])

    evaluator = create_supervised_evaluator(
        model,
        model_transform = model_transform,
        metrics={
            'accuracy': Accuracy(),
            'loss': Loss(loss_fn)},
        prepare_batch=prepare_batch,
        device=params['train']['device'])
```

versus

```python
    trainer = create_supervised_trainer(
        model,
        optimizer,
        model_transform = lambda output: output[0],
        loss_fn=loss_fn,
        prepare_batch=prepare_batch, 
        device=params['train']['device'])

    evaluator = create_supervised_evaluator(
        model,
        # note need to reverse logits and labels!
        output_transform=lambda x, y, y_pred : (y_pred[0], y),
        metrics={
            'accuracy': Accuracy(),
            'loss': Loss(loss_fn)},
        prepare_batch=prepare_batch,
        device=params['train']['device'])
```

But its not a big deal
Thanks @david-waterworth , let's finally add back `model_transform` to `create_supervised_evaluator` as code snippet looks simplier.

Resolving patch:
diff --git a/ignite/engine/__init__.py b/ignite/engine/__init__.py
--- a/ignite/engine/__init__.py
+++ b/ignite/engine/__init__.py
@@ -524,7 +524,7 @@ def output_transform_fn(x, y, y_pred, loss):
     .. versionchanged:: 0.4.7
         Added Gradient Accumulation argument for all supervised training methods.
     .. versionchanged:: 0.4.11
-        Added `model_transform` to transform model's output
+        Added ``model_transform`` to transform model's output
     """
 
     device_type = device.type if isinstance(device, torch.device) else device
@@ -593,6 +593,7 @@ def supervised_evaluation_step(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
+    model_transform: Callable[[Any], Any] = lambda output: output,
     output_transform: Callable[[Any, Any, Any], Any] = lambda x, y, y_pred: (y_pred, y),
 ) -> Callable:
     """
@@ -606,6 +607,8 @@ def supervised_evaluation_step(
             with respect to the host. For other cases, this argument has no effect.
         prepare_batch: function that receives `batch`, `device`, `non_blocking` and outputs
             tuple of tensors `(batch_x, batch_y)`.
+        model_transform: function that receives the output from the model and convert it into the predictions:
+            ``y_pred = model_transform(model(x))``.
         output_transform: function that receives 'x', 'y', 'y_pred' and returns value
             to be assigned to engine's state.output after each iteration. Default is returning `(y_pred, y,)` which fits
             output expected by metrics. If you change it you should use `output_transform` in metrics.
@@ -624,13 +627,16 @@ def supervised_evaluation_step(
         The `model` should be moved by the user before creating an optimizer.
 
     .. versionadded:: 0.4.5
+    .. versionchanged:: 0.4.12
+        Added ``model_transform`` to transform model's output
     """
 
     def evaluate_step(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, Tuple[torch.Tensor]]:
         model.eval()
         with torch.no_grad():
             x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
-            y_pred = model(x)
+            output = model(x)
+            y_pred = model_transform(output)
             return output_transform(x, y, y_pred)
 
     return evaluate_step
@@ -641,6 +647,7 @@ def supervised_evaluation_step_amp(
     device: Optional[Union[str, torch.device]] = None,
     non_blocking: bool = False,
     prepare_batch: Callable = _prepare_batch,
+    model_transform: Callable[[Any], Any] = lambda output: output,
     output_transform: Callable[[Any, Any, Any], Any] = lambda x, y, y_pred: (y_pred, y),
 ) -> Callable:
     """
@@ -654,6 +661,8 @@ def supervised_evaluation_step_amp(
             with respect to the host. For other cases, this argument has no effect.
         prepare_batch: function that receives `batch`, `device`, `non_blocking` and outputs
             tuple of tensors `(batch_x, batch_y)`.
+        model_transform: function that receives the output from the model and convert it into the predictions:
+            ``y_pred = model_transform(model(x))``.
         output_transform: function that receives 'x', 'y', 'y_pred' and returns value
             to be assigned to engine's state.output after each iteration. Default is returning `(y_pred, y,)` which fits
             output expected by metrics. If you change it you should use `output_transform` in metrics.
@@ -672,6 +681,8 @@ def supervised_evaluation_step_amp(
         The `model` should be moved by the user before creating an optimizer.
 
     .. versionadded:: 0.4.5
+    .. versionchanged:: 0.4.12
+        Added ``model_transform`` to transform model's output
     """
     try:
         from torch.cuda.amp import autocast
@@ -683,7 +694,8 @@ def evaluate_step(engine: Engine, batch: Sequence[torch.Tensor]) -> Union[Any, T
         with torch.no_grad():
             x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
             with autocast(enabled=True):
-                y_pred = model(x)
+                output = model(x)
+                y_pred = model_transform(output)
             return output_transform(x, y, y_pred)
 
     return evaluate_step
@@ -711,6 +723,8 @@ def create_supervised_evaluator(
             with respect to the host. For other cases, this argument has no effect.
         prepare_batch: function that receives `batch`, `device`, `non_blocking` and outputs
             tuple of tensors `(batch_x, batch_y)`.
+        model_transform: function that receives the output from the model and convert it into the predictions:
+            ``y_pred = model_transform(model(x))``.
         output_transform: function that receives 'x', 'y', 'y_pred' and returns value
             to be assigned to engine's state.output after each iteration. Default is returning `(y_pred, y,)` which fits
             output expected by metrics. If you change it you should use `output_transform` in metrics.
@@ -737,7 +751,9 @@ def create_supervised_evaluator(
         - `PyTorch's Explanation <https://github.com/pytorch/pytorch/issues/7844#issuecomment-503713840>`_
 
     .. versionchanged:: 0.4.5
-        - Added ``amp_mode`` argument for automatic mixed precision.
+        Added ``amp_mode`` argument for automatic mixed precision.
+    .. versionchanged:: 0.4.12
+        Added ``model_transform`` to transform model's output
     """
     device_type = device.type if isinstance(device, torch.device) else device
     on_tpu = "xla" in device_type if device_type is not None else False
@@ -745,9 +761,23 @@ def create_supervised_evaluator(
 
     metrics = metrics or {}
     if mode == "amp":
-        evaluate_step = supervised_evaluation_step_amp(model, device, non_blocking, prepare_batch, output_transform)
+        evaluate_step = supervised_evaluation_step_amp(
+            model,
+            device,
+            non_blocking=non_blocking,
+            prepare_batch=prepare_batch,
+            model_transform=model_transform,
+            output_transform=output_transform,
+        )
     else:
-        evaluate_step = supervised_evaluation_step(model, device, non_blocking, prepare_batch, output_transform)
+        evaluate_step = supervised_evaluation_step(
+            model,
+            device,
+            non_blocking=non_blocking,
+            prepare_batch=prepare_batch,
+            model_transform=model_transform,
+            output_transform=output_transform,
+        )
 
     evaluator = Engine(evaluate_step)
 


Test patch:
diff --git a/tests/ignite/engine/test_create_supervised.py b/tests/ignite/engine/test_create_supervised.py
--- a/tests/ignite/engine/test_create_supervised.py
+++ b/tests/ignite/engine/test_create_supervised.py
@@ -220,8 +220,18 @@ def _default_create_supervised_evaluator(
     evaluator_device: Optional[str] = None,
     trace: bool = False,
     amp_mode: str = None,
+    with_model_transform: bool = False,
 ):
-    model = DummyModel()
+    if with_model_transform:
+
+        def get_first_element(output):
+            return output[0]
+
+        model = DummyModel(output_as_list=True)
+        model_transform = get_first_element
+    else:
+        model = DummyModel()
+        model_transform = None
 
     if model_device:
         model.to(model_device)
@@ -232,7 +242,12 @@ def _default_create_supervised_evaluator(
         example_input = torch.randn(1, 1)
         model = torch.jit.trace(model, example_input)
 
-    evaluator = create_supervised_evaluator(model, device=evaluator_device, amp_mode=amp_mode)
+    evaluator = create_supervised_evaluator(
+        model,
+        device=evaluator_device,
+        amp_mode=amp_mode,
+        model_transform=model_transform if model_transform is not None else lambda x: x,
+    )
 
     assert model.fc.weight.data[0, 0].item() == approx(0.0)
 
@@ -244,9 +259,14 @@ def _test_create_supervised_evaluator(
     evaluator_device: Optional[str] = None,
     trace: bool = False,
     amp_mode: str = None,
+    with_model_transform: bool = False,
 ):
     model, evaluator = _default_create_supervised_evaluator(
-        model_device=model_device, evaluator_device=evaluator_device, trace=trace, amp_mode=amp_mode
+        model_device=model_device,
+        evaluator_device=evaluator_device,
+        trace=trace,
+        amp_mode=amp_mode,
+        with_model_transform=with_model_transform,
     )
     x = torch.tensor([[1.0], [2.0]])
     y = torch.tensor([[3.0], [5.0]])