make experiment.test return the first generator item

delira-dev · Jun 3, 2019 · 13d3adb · 13d3adb
1 parent b510ffb
commit 13d3adb
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/delira/training/experiment.py b/delira/training/experiment.py
@@ -388,8 +388,9 @@ def test(self, network, test_data: BaseDataManager,
                                convert_batch_to_npy_fn=convert_fn,
                                prepare_batch_fn=prepare_batch, **kwargs)
 
-        return predictor.predict_data_mgr(test_data, 1, metrics,
-                                          metric_keys, verbose)
+        # return first item of generator
+        return next(predictor.predict_data_mgr(test_data, 1, metrics,
+                                               metric_keys, verbose))
 
     def kfold(self, data: BaseDataManager, metrics: dict, num_epochs=None,
               num_splits=None, shuffle=False, random_seed=None,