Repository: pytorch/ignite

Issue:
Improve create_supervised_trainer with optional output_transform
Following [the discussion](https://github.com/pytorch/ignite/pull/476#discussion_r272108999), idea is to give more flexibility to users who are using `create_supervised_trainer`:
```python
def default_output_transform(x, y, y_pred, loss):
    return loss.item() 


def create_supervised_trainer(model, optimizer, loss_fn,
                              device=None, non_blocking=False, prepare_batch=_prepare_batch, 
                              output_transform=default_output_transform):
    if device:
        model.to(device)

    def _update(engine, batch):
        model.train()
        optimizer.zero_grad()
        x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        loss.backward()
        optimizer.step()
        return output_transform(x, y, y_pred, loss)

    return Engine(_update)    
```

cc @IlyaOvodov
Improve create_supervised_trainer with optional output_transform
Following [the discussion](https://github.com/pytorch/ignite/pull/476#discussion_r272108999), idea is to give more flexibility to users who are using `create_supervised_trainer`:
```python
def default_output_transform(x, y, y_pred, loss):
    return loss.item() 


def create_supervised_trainer(model, optimizer, loss_fn,
                              device=None, non_blocking=False, prepare_batch=_prepare_batch, 
                              output_transform=default_output_transform):
    if device:
        model.to(device)

    def _update(engine, batch):
        model.train()
        optimizer.zero_grad()
        x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        loss.backward()
        optimizer.step()
        return output_transform(x, y, y_pred, loss)

    return Engine(_update)    
```

cc @IlyaOvodov


Discussion:
@vfdev-5 
But I think make 2 functions `output_loss` and `output_y_pred_y` and:
```python
def create_supervised_trainer(... , output_transform = output_loss)
def create_supervised_evaluator(... , output_transform = output_y_pred_y)
```
It is more meaningful, IMHO
I'll make it
@vfdev-5 
But I think make 2 functions `output_loss` and `output_y_pred_y` and:
```python
def create_supervised_trainer(... , output_transform = output_loss)
def create_supervised_evaluator(... , output_transform = output_y_pred_y)
```
It is more meaningful, IMHO
I'll make it

Resolving patch:
diff --git a/ignite/engine/__init__.py b/ignite/engine/__init__.py
--- a/ignite/engine/__init__.py
+++ b/ignite/engine/__init__.py
@@ -15,7 +15,8 @@ def _prepare_batch(batch, device=None, non_blocking=False):
 
 def create_supervised_trainer(model, optimizer, loss_fn,
                               device=None, non_blocking=False,
-                              prepare_batch=_prepare_batch):
+                              prepare_batch=_prepare_batch,
+                              output_transform=lambda x, y, y_pred, loss: loss.item()):
     """
     Factory function for creating a trainer for supervised models.
 
@@ -29,8 +30,11 @@ def create_supervised_trainer(model, optimizer, loss_fn,
             with respect to the host. For other cases, this argument has no effect.
         prepare_batch (callable, optional): function that receives `batch`, `device`, `non_blocking` and outputs
             tuple of tensors `(batch_x, batch_y)`.
+        output_transform (callable, optional): function that receives 'x', 'y', 'y_pred', 'loss' and returns value
+            to be assigned to engine's state.output after each iteration. Default is returning `loss.item()`.
 
-    Note: `engine.state.output` for this engine is the loss of the processed batch.
+    Note: `engine.state.output` for this engine is defind by `output_transform` parameter and is the loss
+        of the processed batch by default.
 
     Returns:
         Engine: a trainer engine with supervised update function.
@@ -46,14 +50,15 @@ def _update(engine, batch):
         loss = loss_fn(y_pred, y)
         loss.backward()
         optimizer.step()
-        return loss.item()
+        return output_transform(x, y, y_pred, loss)
 
     return Engine(_update)
 
 
 def create_supervised_evaluator(model, metrics={},
                                 device=None, non_blocking=False,
-                                prepare_batch=_prepare_batch):
+                                prepare_batch=_prepare_batch,
+                                output_transform=lambda x, y, y_pred: (y_pred, y,)):
     """
     Factory function for creating an evaluator for supervised models.
 
@@ -66,8 +71,12 @@ def create_supervised_evaluator(model, metrics={},
             with respect to the host. For other cases, this argument has no effect.
         prepare_batch (callable, optional): function that receives `batch`, `device`, `non_blocking` and outputs
             tuple of tensors `(batch_x, batch_y)`.
+        output_transform (callable, optional): 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.
 
-    Note: `engine.state.output` for this engine is a tuple of `(batch_pred, batch_y)`.
+    Note: `engine.state.output` for this engine is defind by `output_transform` parameter and is
+        a tuple of `(batch_pred, batch_y)` by default.
 
     Returns:
         Engine: an evaluator engine with supervised inference function.
@@ -80,7 +89,7 @@ def _inference(engine, batch):
         with torch.no_grad():
             x, y = prepare_batch(batch, device=device, non_blocking=non_blocking)
             y_pred = model(x)
-            return y_pred, y
+            return output_transform(x, y, y_pred)
 
     engine = Engine(_inference)
 


Test patch: