Handle NER usecase (#263)

* Handle NER usecase

* Fix test

baal/transformers_trainer_wrapper.py

@@ -79,7 +79,7 @@ def predict_on_dataset_generator(
             )
 
             out = map_on_tensor(lambda o: o.view([iterations, -1, *o.size()[1:]]), out)
-            out = map_on_tensor(lambda o: o.permute(1, 2, *range(3, o.ndimension()), 0), out)
+            out = map_on_tensor(lambda o: o.permute(1, *range(3, o.ndimension()), 2, 0), out)
             out = map_on_tensor(lambda x: x.detach(), out)
             if half:
                 out = map_on_tensor(lambda x: x.half(), out)

baal/utils/iterutils.py

@@ -1,10 +1,10 @@
-from collections.abc import Sequence
+from collections.abc import Sequence, MutableMapping
 
 
 def map_on_tensor(fn, val):
     """Map a function on a Tensor or a list of Tensors"""
     if isinstance(val, Sequence):
-        return [fn(v) for v in val]
-    elif isinstance(val, dict):
-        return {k: fn(v) for k, v in val.items()}
+        return [map_on_tensor(fn, v) for v in val]
+    elif isinstance(val, (dict, MutableMapping)):
+        return {k: map_on_tensor(fn, v) for k, v in val.items()}
     return fn(val)

tests/transformers_trainer_wrapper_test.py

@@ -40,11 +40,11 @@ def test_predict_on_dataset_generator(self):
 
         # iteration == 1
         pred = self.wrapper.predict_on_dataset_generator(self.dataset, 1, False)
-        assert next(pred).shape == (5, 10, 100, 1)
+        assert next(pred).shape == (5, 100, 10, 1)
 
         # iterations > 1
         pred = self.wrapper.predict_on_dataset_generator(self.dataset, 10, False)
-        assert next(pred).shape == (5, 10, 100, 10)
+        assert next(pred).shape == (5, 100, 10, 10)
 
         # Test generators
         l_gen = self.wrapper.predict_on_dataset_generator(self.dataset, 10, False)