diff --git a/mushroom_rl/core/dataset.py b/mushroom_rl/core/dataset.py
index 01879c65..919d7900 100644
--- a/mushroom_rl/core/dataset.py
+++ b/mushroom_rl/core/dataset.py
@@ -291,7 +291,7 @@ def episodes_length(self):
                 lengths.append(l)
                 l = 0
 
-        return lengths
+        return self._array_backend.from_list(lengths)
 
     @property
     def n_episodes(self):