Repository: pytorch/ignite

Issue:
Match core API in contrib TBPTT trainer
I propose the following change so that tbptt trainer behaves like the supervised trainer:

- [ ] Replace `convert_tensor` by `prepare_batch` ([line](https://github.com/pytorch/ignite/blob/d3d7b146d56efb5b497f3cf4b22caab976774e7e/ignite/contrib/engines/tbptt.py#L29))
- [ ] Call `prepare_batch` on time chunks (the quantity used to compute gradients) rather that whole time series mini bactch
- [ ] Add `prepare_batch` function in the arguments
- [ ] Add `non_blocking` in the arguments
Match core API in contrib TBPTT trainer
I propose the following change so that tbptt trainer behaves like the supervised trainer:

- [ ] Replace `convert_tensor` by `prepare_batch` ([line](https://github.com/pytorch/ignite/blob/d3d7b146d56efb5b497f3cf4b22caab976774e7e/ignite/contrib/engines/tbptt.py#L29))
- [ ] Call `prepare_batch` on time chunks (the quantity used to compute gradients) rather that whole time series mini bactch
- [ ] Add `prepare_batch` function in the arguments
- [ ] Add `non_blocking` in the arguments


Discussion:

Resolving patch:
diff --git a/ignite/contrib/engines/tbptt.py b/ignite/contrib/engines/tbptt.py
--- a/ignite/contrib/engines/tbptt.py
+++ b/ignite/contrib/engines/tbptt.py
@@ -4,8 +4,8 @@
 
 import torch
 
-from ignite._utils import convert_tensor, apply_to_tensor
-from ignite.engine import Engine
+from ignite._utils import apply_to_tensor
+from ignite.engine import Engine, _prepare_batch
 
 
 class Tbptt_Events(Enum):
@@ -19,18 +19,6 @@ class Tbptt_Events(Enum):
     TIME_ITERATION_COMPLETED = "time_iteration_completed"
 
 
-def _prepare_tbptt_batch(batch, tbptt_step, dim=0, device=None):
-    """Prepare batch for tbptt trainer.
-
-    Batch come from the dataloader. It is split in chunks along the time
-    dimension and fed to the truncated backpropagation throught time trainer.
-    """
-    x, y = batch
-    x = convert_tensor(x, device=device)
-    y = convert_tensor(y, device=device)
-    return zip(x.split(tbptt_step, dim=dim), y.split(tbptt_step, dim=dim))
-
-
 def _detach_hidden(hidden):
     """Cut backpropagation graph.
 
@@ -46,7 +34,9 @@ def create_supervised_tbptt_trainer(
     loss_fn,
     tbtt_step,
     dim=0,
-    device=None
+    device=None,
+    non_blocking=False,
+    prepare_batch=_prepare_batch
 ):
     """Create a trainer for truncated backprop through time supervised models.
 
@@ -70,6 +60,11 @@ def create_supervised_tbptt_trainer(
         dim (int): axis representing the time dimension
         device (str, optional): device type specification (default: None).
             Applies to both model and batches.
+        non_blocking (bool, optional): if True and this copy is between CPU and GPU,
+            the copy may occur asynchronously 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)`.
 
     Returns:
         Engine: a trainer engine with supervised update function
@@ -82,11 +77,9 @@ def _update(engine, batch):
         loss_list = []
         hidden = None
 
-        # Batches split in time chunks
-        batch_splits = _prepare_tbptt_batch(
-            batch, tbtt_step, dim=dim, device=device
-        )
-        for x_t, y_t in batch_splits:
+        x, y = batch
+        for batch_t in zip(x.split(tbtt_step, dim=dim), y.split(tbtt_step, dim=dim)):
+            x_t, y_t = prepare_batch(batch_t, device=device, non_blocking=non_blocking)
             # Fire event for start of iteration
             engine.fire_event(Tbptt_Events.TIME_ITERATION_STARTED)
             # Forward, backward and
diff --git a/ignite/engine/__init__.py b/ignite/engine/__init__.py
--- a/ignite/engine/__init__.py
+++ b/ignite/engine/__init__.py
@@ -5,6 +5,9 @@
 
 
 def _prepare_batch(batch, device=None, non_blocking=False):
+    """Prepare batch for training: pass to a device with options
+
+    """
     x, y = batch
     return (convert_tensor(x, device=device, non_blocking=non_blocking),
             convert_tensor(y, device=device, non_blocking=non_blocking))


Test patch: