diff --git a/docs/tutorial/library.md b/docs/tutorial/library.md
index 0153309..e951ec2 100644
--- a/docs/tutorial/library.md
+++ b/docs/tutorial/library.md
@@ -12,38 +12,61 @@ if os.path.exists('examples/mlruns'):
     shutil.rmtree('examples/mlruns')
 ```
 
-## PyTorch vs TensorFlow
+## Base Parameter File
 
-In this section we compare two libraries and show that using different libraries on the same problem is straightforward.
+Before examples, we write two base or *template* parameter files which are extened by other parameter files.
+
+#File A base parameter YAML file for various libraries (data.yml) {%=/examples/data.yml%}
+
+#File A base parameter YAML file for various libraries (base.yml) {%=/examples/base.yml%} {#base#}
+
+In `base.yml`, the first line "`extends: data`" means that the file extends (or includes, in this case) `data.yml`.
+
+## Neural Network Libraries
+
+In this section we compare three libraries ([TensorFlow](https://www.tensorflow.org/), [NNabla](https://nnabla.org/), and [PyTorch](https://pytorch.org/)), and show that using different libraries on the same problem is straightforward.
+
+```python
+import tensorflow
+import nnabla
+import torch
+print(tensorflow.__version__)
+print(nnabla.__version__)
+print(torch.__version__)
+```
 
 First define models:
 
-#File A Model definition in PyTorch (rectangle/torch.py) {%=/examples/rectangle/torch.py%}
+#File A Model definition in TensorFlow (rectangle/tensorflow.py) {%=/examples/rectangle/tensorflow.py%}
 
-#File A Model definition in TensorFlow (rectangle/tf.py) {%=/examples/rectangle/tf.py%}
+#File A Model definition in NNabla (rectangle/nnabla.py) {%=/examples/rectangle/nnabla.py%}
+
+#File A Model definition in PyTorch (rectangle/torch.py) {%=/examples/rectangle/torch.py%}
 
 For simplicity, the TensorFlow model is defined by using the `keras.Sequential()`, so that we call the `create_model()` function to get the model.
 
 Next, write parameter YAML files:
 
-#File A parameter YAML for PyTorch (torch.yml) {%=/examples/torch.yml%}
+#File A parameter YAML file for TensorFlow (tensorflow.yml) {%=/examples/tensorflow.yml%}
 
-#File A parameter YAML for TensorFlow (tf.yml) {%=/examples/tf.yml%}
+#File A parameter YAML file for NNabla (nnabla.yml) {%=/examples/nnabla.yml%}
 
-These YAML files have a very similar structure. The first difference comes from that, in PyTorch, an optimizer needs model parameters at the time of instantiation and a scheduler needs an optimizer too, while, in TensorFlow, an optimizer and scheduler can be instantiated without other related instances.  The second difference is `loss` functions. The Pytorch YAML file writes the fullname, while the TensorFlow one writes an abbreviation.
+#File A parameter YAML fine for PyTorch (torch2.yml) {%=/examples/torch2.yml%}
 
+These YAML files are very similar. The only difference is that, in PyTorch, an optimizer needs model parameters at the time of instantiation.
 
 !!! note
     The `model` for TensorFlow is a function. A new `call` key is used. (But you can use `class`, too, or `call` for a class, vice versa, because both a class and function are *callable*.)
 
-Next, create two runs.
+Next, create three runs.
 
 ```python
 import ivory
 
 client = ivory.create_client("examples")
-run_torch = client.create_run('torch')
-run_tf = client.create_run('tf')
+run_tf = client.create_run('tensorflow')
+run_nn = client.create_run('nnabla')
+run_torch = client.create_run('torch2')
 ```
 
 For comparison, equalize initial parameters.
@@ -51,21 +74,34 @@ For comparison, equalize initial parameters.
 ```python
 import torch
 
-for w_tf, w_torch in zip(run_tf.model.weights, run_torch.model.parameters()):
+# These three lines are only needed for this example.
+run, trainer = run_nn, run_nn.trainer
+run.model.build(trainer.loss, run.datasets.train, trainer.batch_size)
+run.optimizer.set_parameters(run.model.parameters())
+
+ws_tf = run_tf.model.weights
+ws_nn = run_nn.model.parameters().values()
+ws_torch = run_torch.model.parameters()
+for w_tf, w_nn, w_torch in zip(ws_tf, ws_nn, ws_torch):
+    w_nn.data.data = w_tf.numpy()
     w_torch.data = torch.tensor(w_tf.numpy().T)
 ```
 
 Then, start the runs.
 
 ```python
-run_torch.start()
+run_tf.start('both')  # Slower due to usage of GPU for a simple network.
+```
+
+```python
+run_nn.start('both')
 ```
 
 ```python
-run_tf.start()
+run_torch.start('both')
 ```
 
-Visualize the results:
+Metrics during training are almost same. Visualize the results:
 
 ```python
 import matplotlib.pyplot as plt
@@ -79,28 +115,47 @@ def plot(run):
     plt.xlabel('Target values')
     plt.ylabel('Predicted values')
 
-for run in [run_tf, run_torch]:
+for run in [run_tf, run_nn, run_torch]:
     plot(run)
 ```
 
 Actual outputs are like below:
 
 ```python
-x = run_tf.datasets.test[:10][1]
+x = run_tf.datasets.val[:5][1]
 run_tf.model.predict(x)
 ```
 
 ```python
+x = run_nn.datasets.val[:5][1]
+run_nn.model(x)
+```
+
+```python
+x = run_torch.datasets.val[:5][1]
 run_torch.model(torch.tensor(x))
 ```
 
+You can *ensemble* these resutls, although this is meaningless in this example.
+
+```python
+from ivory.callbacks.results import concatenate
+
+results = concatenate(list(run.results for run in [run_tf, run_nn, run_torch]))
+index = results.val.index.argsort()
+results.val.output[index[:15]]
+```
+
+```python
+reduced_results = results.mean()
+reduced_results.val.output[:5]
+```
+
 ## Scikit-learn
 
 Ivory can optimize various scikit-learn's estimators. Here are som examples.
 
-### RandomForestRegressor
-
-#File A parameter YAML for RandomForestRegressor (rfr.yml) {%=/examples/rfr.yml%}
+#File A base parameter YAML file for various estimators (data2.yml) {%=/examples/data2.yml%}
 
 The `dataset` has a `transform` argument. This function reshapes the target array to match the shape for scikit-learn estimators (1D from 2D).
 
@@ -108,6 +163,10 @@ The `dataset` has a `transform` argument. This function reshapes the target arra
 
 #Code rectangle.data.transform() {{ rectangle.data.transform # inspect }}
 
+### RandomForestRegressor
+
+#File A parameter YAML file for RandomForestRegressor (rfr.yml) {%=/examples/rfr.yml%}
+
 There are nothing different to start a run.
 
 ```python
@@ -123,7 +182,7 @@ plot(run)
 
 ### Ridge
 
-#File A parameter YAML for Ridge (ridge.yml) {%=/examples/ridge.yml%}
+#File A parameter YAML file for Ridge (ridge.yml) {%=/examples/ridge.yml%}
 
 Because the Ridge estimator has no `criterion` attribute, you have to specify metrics if you need. A `mse` key has empty (`None`) value. In this case, the default function (`sklearn.metrics.mean_squared_error()`) is chosen. On the other hand, `mse_2`'s value is a custom funtion's name:
 
@@ -149,7 +208,7 @@ For LightGBM, Ivory implements two estimators:
 * `ivory.lightgbm.estimator.Regressor`
 * `ivory.lightgbm.estimator.Classifier`
 
-#File A parameter YAML for LightGBM (lgb.yml) {%=/examples/lgb.yml%}
+#File A parameter YAML file for LightGBM (lgb.yml) {%=/examples/lgb.yml%}
 
 ```python
 run = client.create_run('lgb')
diff --git a/examples/base.yml b/examples/base.yml
new file mode 100644
index 0000000..16de3ef
--- /dev/null
+++ b/examples/base.yml
@@ -0,0 +1,15 @@
+extends: data
+model:
+  hidden_sizes: [20, 30]
+optimizer:
+  lr: 1e-3
+results:
+metrics:
+monitor:
+  metric: val_loss
+trainer:
+  loss: mse
+  batch_size: 5
+  epochs: 5
+  shuffle: false
+  verbose: 2
diff --git a/examples/data.yml b/examples/data.yml
new file mode 100644
index 0000000..3c00912
--- /dev/null
+++ b/examples/data.yml
@@ -0,0 +1,6 @@
+datasets:
+  data:
+    class: rectangle.data.Data
+    n_splits: 4
+  dataset:
+  fold: 0
diff --git a/examples/data2.yml b/examples/data2.yml
new file mode 100644
index 0000000..dfff1fd
--- /dev/null
+++ b/examples/data2.yml
@@ -0,0 +1,4 @@
+extends: data
+datasets:
+  dataset:
+    transform: rectangle.data.transform
diff --git a/examples/lgb.yml b/examples/lgb.yml
index 3ff630c..3ee1488 100644
--- a/examples/lgb.yml
+++ b/examples/lgb.yml
@@ -1,10 +1,4 @@
-datasets:
-  data:
-    class: rectangle.data.Data
-    n_splits: 4
-  dataset:
-    transform: rectangle.data.transform
-  fold: 0
+extends: data2
 estimator:
   class: ivory.lightgbm.estimator.Regressor
   boosting_type: gbdt
diff --git a/examples/nnabla.yml b/examples/nnabla.yml
index e444fd0..2809526 100644
--- a/examples/nnabla.yml
+++ b/examples/nnabla.yml
@@ -1,25 +1,6 @@
 library: nnabla
-datasets:
-  data:
-    class: rectangle.data.Data
-    n_splits: 4
-  dataset:
-  fold: 0
+extends: base
 model:
   class: rectangle.nnabla.Model
-  hidden_sizes: [20, 30]
 optimizer:
   class: nnabla.solvers.Sgd
-  lr: 1e-3
-results:
-metrics:
-monitor:
-  metric: val_loss
-early_stopping:
-  patience: 10
-trainer:
-  loss: mse
-  batch_size: 10
-  epochs: 10
-  shuffle: true
-  verbose: 2
diff --git a/examples/rectangle/tf.py b/examples/rectangle/tensorflow.py
similarity index 100%
rename from examples/rectangle/tf.py
rename to examples/rectangle/tensorflow.py
diff --git a/examples/rfr.yml b/examples/rfr.yml
index ac22a89..585f963 100644
--- a/examples/rfr.yml
+++ b/examples/rfr.yml
@@ -1,11 +1,5 @@
 library: sklearn
-datasets:
-  data:
-    class: rectangle.data.Data
-    n_splits: 4
-  dataset:
-    transform: rectangle.data.transform
-  fold: 0
+extends: data2
 estimator:
   model: sklearn.ensemble.RandomForestRegressor
   n_estimators: 5
diff --git a/examples/ridge.yml b/examples/ridge.yml
index e3b1b3b..a2a0e1f 100644
--- a/examples/ridge.yml
+++ b/examples/ridge.yml
@@ -1,11 +1,5 @@
 library: sklearn
-datasets:
-  data:
-    class: rectangle.data.Data
-    n_splits: 4
-  dataset:
-    transform: rectangle.data.transform
-  fold: 0
+extends: data2
 estimator:
   model: sklearn.linear_model.Ridge
 results:
diff --git a/examples/tensorflow.yml b/examples/tensorflow.yml
new file mode 100644
index 0000000..2939103
--- /dev/null
+++ b/examples/tensorflow.yml
@@ -0,0 +1,6 @@
+library: tensorflow
+extends: base
+model:
+  call: rectangle.tensorflow.create_model
+optimizer:
+  class: tensorflow.keras.optimizers.SGD
diff --git a/examples/tf.yml b/examples/tf.yml
deleted file mode 100644
index ae90768..0000000
--- a/examples/tf.yml
+++ /dev/null
@@ -1,29 +0,0 @@
-library: tensorflow
-datasets:
-  data:
-    class: rectangle.data.Data
-    n_splits: 4
-  dataset:
-  fold: 0
-model:
-  call: rectangle.tf.create_model
-  hidden_sizes: [20, 30]
-optimizer:
-  class: tensorflow.keras.optimizers.SGD
-  lr: 1e-3
-scheduler:
-  class: tensorflow.keras.callbacks.ReduceLROnPlateau
-  factor: 0.5
-  patience: 4
-results:
-metrics:
-monitor:
-  metric: val_loss
-early_stopping:
-  patience: 10
-trainer:
-  loss: mse
-  batch_size: 10
-  epochs: 10
-  shuffle: true
-  verbose: 2
diff --git a/examples/torch2.yml b/examples/torch2.yml
new file mode 100644
index 0000000..089a6f4
--- /dev/null
+++ b/examples/torch2.yml
@@ -0,0 +1,7 @@
+library: torch
+extends: base
+model:
+  class: rectangle.torch.Model
+optimizer:
+  class: torch.optim.SGD
+  _: $.model.parameters()
diff --git a/ivory/__init__.py b/ivory/__init__.py
index c69d7a2..f2a3053 100644
--- a/ivory/__init__.py
+++ b/ivory/__init__.py
@@ -1,4 +1,4 @@
-__version__ = "0.3.0"
+__version__ = "0.3.1"
 
 from ivory.core.client import create_client
 
diff --git a/ivory/callbacks/results.py b/ivory/callbacks/results.py
index 2fef0d3..ca4cbe5 100644
--- a/ivory/callbacks/results.py
+++ b/ivory/callbacks/results.py
@@ -1,5 +1,5 @@
 """A container to store training, validation and test results. """
-from typing import Callable, Dict, Iterable, List, Optional
+from typing import Callable, Dict, Iterable, Optional
 
 import numpy as np
 import pandas as pd
@@ -42,16 +42,16 @@ def result_dict(self):
         dict = ivory.core.collections.Dict()
         return dict(index=self.index, output=self.output, target=self.target)
 
-    def set(self, **kwargs):
-        results = {}
-        for key, value in kwargs.items():
-            dict = ivory.core.collections.Dict()
-            if len(value) == 3:
-                dict(index=value[0], output=value[1], target=value[2])
-            else:
-                dict(index=value[0], output=value[1], target=None)
-            results[key] = dict
-        super().set(**results)
+    # def set(self, **kwargs):
+    #     results = {}
+    #     for key, value in kwargs.items():
+    #         dict = ivory.core.collections.Dict()
+    #         if len(value) == 3:
+    #             dict(index=value[0], output=value[1], target=value[2])
+    #         else:
+    #             dict(index=value[0], output=value[1], target=None)
+    #         results[key] = dict
+    #     super().set(**results)
 
     def mean(self):
         results = Results()
@@ -98,11 +98,11 @@ def result_dict(self):
         return super().result_dict()
 
 
-def stack(x: List[np.ndarray]) -> np.ndarray:
-    if x[0].ndim == 1:
-        return np.hstack(x)
-    else:
-        return np.vstack(x)
+# def stack(x: List[np.ndarray]) -> np.ndarray:
+#     if x[0].ndim == 1:
+#         return np.hstack(x)
+#     else:
+#         return np.vstack(x)
 
 
 def concatenate(
diff --git a/ivory/core/data.py b/ivory/core/data.py
index dea1cf4..5596787 100644
--- a/ivory/core/data.py
+++ b/ivory/core/data.py
@@ -24,9 +24,7 @@
     Use a `'def'` key for `dataset` instead of `'class'`.
     See [Tutorial](/tutorial/data)
 """
-
-
-from dataclasses import InitVar, dataclass
+from dataclasses import dataclass
 from typing import Callable, Optional, Tuple
 
 import numpy as np
diff --git a/ivory/core/trainer.py b/ivory/core/trainer.py
index 993fb2b..ccdbb5a 100644
--- a/ivory/core/trainer.py
+++ b/ivory/core/trainer.py
@@ -3,6 +3,7 @@
 from optuna.exceptions import TrialPruned
 from termcolor import colored
 
+from ivory.core import instance
 from ivory.core.exceptions import EarlyStopped, Pruned
 from ivory.core.run import Run
 from ivory.core.state import State
@@ -14,6 +15,9 @@ class Trainer(State):
     epoch: int = -1
     epochs: int = 1
     global_step: int = -1
+    batch_size: int = 32
+    shuffle: bool = True
+    dataloaders: str = ""
     verbose: int = 1
 
     def start(self, run: Run):
@@ -27,6 +31,14 @@ def start(self, run: Run):
         else:
             self.test(run)
 
+    def on_init_begin(self, run):
+        if not run.dataloaders:
+            dataloaders_factory = instance.get_attr(self.dataloaders)
+            dataloaders = dataloaders_factory(
+                run.datasets, self.batch_size, self.shuffle
+            )
+            run.set(dataloaders=dataloaders)
+
     def train(self, run: Run):
         run.on_fit_begin()
         try:
diff --git a/ivory/nnabla/data.py b/ivory/nnabla/data.py
index a58415f..ad38f30 100644
--- a/ivory/nnabla/data.py
+++ b/ivory/nnabla/data.py
@@ -25,6 +25,8 @@ def __post_init__(self):
         )
 
     def __len__(self):
+        if len(self.dataset) % self.batch_size:  # FIXME
+            raise NotImplementedError
         return len(self.dataset) // self.batch_size
 
     def load_func(self, index):
diff --git a/ivory/nnabla/trainer.py b/ivory/nnabla/trainer.py
index cfe7a25..8957903 100644
--- a/ivory/nnabla/trainer.py
+++ b/ivory/nnabla/trainer.py
@@ -9,14 +9,12 @@
 import ivory.nnabla.data
 import ivory.nnabla.functions
 from ivory.core import instance
-from ivory.core.run import Run
 
 
 @dataclass
 class Trainer(ivory.core.trainer.Trainer):
     loss: Optional[Callable] = None
-    batch_size: int = 32
-    shuffle: bool = True
+    dataloaders: str = "ivory.nnabla.data.DataLoaders"
     gpu: bool = False
     precision: int = 32  # Full precision (32), half precision (16).
     amp_level: str = "O1"
@@ -28,6 +26,7 @@ def __post_init__(self):
             self.loss = instance.get_attr(self.loss)
 
     def on_init_begin(self, run):
+        super().on_init_begin(run)
         if self.gpu:
             context = "cudnn"
         else:
@@ -46,12 +45,6 @@ def on_init_begin(self, run):
             run.model.build(self.loss, run.datasets.train, self.batch_size)
             run.optimizer.set_parameters(run.model.parameters())
 
-        if not run.dataloaders:
-            dataloaders = ivory.nnabla.data.DataLoaders(
-                run.datasets, self.batch_size, self.shuffle
-            )
-            run.set(dataloaders=dataloaders)
-
     def on_train_begin(self, run):
         run.model.train()
 
@@ -78,6 +71,6 @@ def on_epoch_end(self, run):
     def on_test_begin(self, run):
         run.model.eval()
 
-    def test_step(self, run, index, input, *target):
+    def test_step(self, run, index, input, target):
         output = run.model(input)
         run.results.step(index, output, target)
diff --git a/ivory/torch/trainer.py b/ivory/torch/trainer.py
index c8baf0b..9bde3ad 100644
--- a/ivory/torch/trainer.py
+++ b/ivory/torch/trainer.py
@@ -20,8 +20,7 @@
 @dataclass
 class Trainer(ivory.core.trainer.Trainer):
     loss: Optional[Callable] = None
-    batch_size: int = 32
-    shuffle: bool = True
+    dataloaders: str = "ivory.torch.data.DataLoaders"
     gpu: bool = False
     precision: int = 32  # Full precision (32), half precision (16).
     amp_level: str = "O1"
@@ -34,17 +33,13 @@ def __post_init__(self):
             self.loss = instance.get_attr(self.loss)
 
     def on_init_begin(self, run):
+        super().on_init_begin(run)
         if self.gpu:
             run.model.cuda()
             if self.precision == 16:
                 run.model, run.optimizer = amp.initialize(
                     run.model, run.optimizer, opt_level=self.amp_level
                 )
-        if not run.dataloaders:
-            dataloaders = ivory.torch.data.DataLoaders(
-                run.datasets, self.batch_size, self.shuffle
-            )
-            run.set(dataloaders=dataloaders)
 
     def on_train_begin(self, run):
         run.model.train()
diff --git a/ivory/utils/path.py b/ivory/utils/path.py
index d29fe3d..d523c6e 100644
--- a/ivory/utils/path.py
+++ b/ivory/utils/path.py
@@ -87,7 +87,7 @@ def update_include(params, source_name, include=None):
 
 
 def inherit(params, source_name):
-    if "base" in params:
+    if "extends" in params:
         return _inherit(params, source_name)
     for key, value in params.items():
         if isinstance(value, dict):
@@ -96,8 +96,13 @@ def inherit(params, source_name):
 
 
 def _inherit(params, source_name):
-    name = params.pop("base")
-    base = load_params(name, source_name)[0]
+    base = {}
+    for key, value in params.items():
+        if key == 'extends':
+            break
+        base[key] = value
+    name = params.pop("extends")
+    base.update(load_params(name, source_name)[0])
     for key, value in base.items():
         if key in params:
             if value is None:
@@ -105,6 +110,9 @@ def _inherit(params, source_name):
             elif isinstance(value, dict):
                 for k in params[key]:
                     value[k] = params[key][k]
+    for key, value in params.items():
+        if key not in base:
+            base[key] = value
     return base
 
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 3aade12..e79d810 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -3,6 +3,7 @@
 import sys
 
 import pytest
+import tensorflow as tf
 
 from ivory.core.client import create_client
 
@@ -10,7 +11,15 @@
 
 
 @pytest.fixture(scope="session")
-def client():
+def setup():
+    gpus = tf.config.experimental.list_physical_devices("GPU")
+    if gpus:
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+
+
+@pytest.fixture(scope="session")
+def client(setup):
     client = create_client(directory="tests/examples")
     yield client
     if os.path.exists("tests/examples/mlruns"):
diff --git a/tests/examples/example.yaml b/tests/examples/example.yaml
index 71242c2..255e0ad 100644
--- a/tests/examples/example.yaml
+++ b/tests/examples/example.yaml
@@ -8,7 +8,7 @@ study:
       def: example.suggest_hidden_sizes
       max_num_layers: 3
 run:
-  base: base
+  extends: base
   dataloaders:
     data:
       class: example.TorchData
diff --git a/tests/examples/lgbr.yaml b/tests/examples/lgbr.yaml
index 1241732..7e7d7d8 100644
--- a/tests/examples/lgbr.yaml
+++ b/tests/examples/lgbr.yaml
@@ -12,3 +12,4 @@ estimator:
   num_boost_round: 10
 results:
 metrics:
+  mse:
diff --git a/tests/libs/conftest.py b/tests/libs/conftest.py
new file mode 100644
index 0000000..28a2049
--- /dev/null
+++ b/tests/libs/conftest.py
@@ -0,0 +1,36 @@
+import os
+import shutil
+import sys
+
+import pytest
+import torch
+
+from ivory.core.client import create_client
+
+
+@pytest.fixture(scope="module")
+def runs():
+    sys.path.insert(0, os.path.abspath("examples"))
+    client = create_client(directory="examples")
+    runs = []
+    for name in ["tensorflow", "nnabla", "torch2"]:
+        run = client.create_run(name, epochs=5, batch_size=10, shuffle=False)
+        runs.append(run)
+    run_tf, run_nn, run_torch = runs
+
+    run_nn.model.build(
+        run_nn.trainer.loss, run_nn.datasets.train, run_nn.trainer.batch_size
+    )
+    run_nn.optimizer.set_parameters(run_nn.model.parameters())
+
+    ws_tf = run_tf.model.weights
+    ws_nn = run_nn.model.parameters().values()
+    ws_torch = run_torch.model.parameters()
+    for w_tf, w_nn, w_torch in zip(ws_tf, ws_nn, ws_torch):
+        w_nn.data.data = w_tf.numpy()
+        w_torch.data = torch.tensor(w_tf.numpy().T)
+
+    yield dict(zip(["tf", "nn", "torch"], runs))
+    del sys.path[0]
+    if os.path.exists("examples/mlruns"):
+        shutil.rmtree("examples/mlruns")
diff --git a/tests/libs/test_libraries.py b/tests/libs/test_libraries.py
new file mode 100644
index 0000000..3362d46
--- /dev/null
+++ b/tests/libs/test_libraries.py
@@ -0,0 +1,23 @@
+import numpy as np
+
+from ivory.callbacks.results import concatenate
+
+
+def test_libraries(runs):
+    for run in runs.values():
+        run.start("both")
+
+    for mode in ["val", "test"]:
+        outputs = []
+        for run in runs.values():
+            outputs.append(run.results[mode].output)
+
+        for output in outputs[1:]:
+            assert np.allclose(outputs[0], output)
+
+    def callback(index, output, target):
+        return index, 2 * output, target
+
+    gen = (run.results for run in runs.values())
+    results = concatenate(gen, reduction="mean", callback=callback)
+    assert np.allclose(2 * outputs[0], results.test.output)
diff --git a/tests/tensorflow/test_tensorflow_mnist.py b/tests/tensorflow/test_tensorflow_mnist.py
index 1258fbc..7bdfb44 100644
--- a/tests/tensorflow/test_tensorflow_mnist.py
+++ b/tests/tensorflow/test_tensorflow_mnist.py
@@ -1,12 +1,7 @@
 import numpy as np
-import tensorflow as tf
 
 
 def test_mnist(client):
-    gpus = tf.config.experimental.list_physical_devices("GPU")
-    if gpus:
-        for gpu in gpus:
-            tf.config.experimental.set_memory_growth(gpu, True)
     run = client.create_run("mnist")
     run.start("train")
     run.start("test")