Fixed packed sequences for numpy

MushroomRL · Dec 6, 2023 · e5de4fb · e5de4fb
1 parent a5dfa3a
commit e5de4fb
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/mushroom_rl/core/_impl/array_backend.py b/mushroom_rl/core/_impl/array_backend.py
@@ -98,8 +98,8 @@ def pack_padded_sequence(array, lengths):
         shape = array.shape
 
         new_shape = (shape[0] * shape[1],) + shape[2:]
-        mask = (np.arange(len(array))[:, None] < lengths[None, :]).flatten()
-        return array.reshape(new_shape)[mask]
+        mask = (np.arange(len(array))[:, None] < lengths[None, :]).flatten(order='F')
+        return array.reshape(new_shape, order='F')[mask]
 
 
 class TorchBackend(ArrayBackend):