Merge branch 'develop' into toni/truncated_signal_not_treated

Toni-SM · Jan 16, 2025 · c6ece72 · c6ece72
2 parents f44465a + 7f9992b
commit c6ece72
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/skrl/memories/jax/base.py b/skrl/memories/jax/base.py
@@ -188,7 +188,8 @@ def create_tensor(
         :rtype: bool
         """
         # compute data size
-        size = compute_space_size(size, occupied_size=True)
+        if not keep_dimensions:
+            size = compute_space_size(size, occupied_size=True)
         # check dtype and size if the tensor exists
         if name in self.tensors:
             tensor = self.tensors[name]
@@ -406,7 +407,7 @@ def sample_all(
         # sequential order
         if sequence_length > 1:
             if mini_batches > 1:
-                batches = np.array_split(self.all_sequence_indexes, len(self.all_sequence_indexes) // mini_batches)
+                batches = np.array_split(self.all_sequence_indexes, mini_batches)
                 return [[self._get_tensors_view(name)[batch] for name in names] for batch in batches]
             return [[self._get_tensors_view(name)[self.all_sequence_indexes] for name in names]]
 

diff --git a/skrl/memories/torch/base.py b/skrl/memories/torch/base.py
@@ -157,7 +157,8 @@ def create_tensor(
         :rtype: bool
         """
         # compute data size
-        size = compute_space_size(size, occupied_size=True)
+        if not keep_dimensions:
+            size = compute_space_size(size, occupied_size=True)
         # check dtype and size if the tensor exists
         if name in self.tensors:
             tensor = self.tensors[name]
@@ -342,7 +343,7 @@ def sample_all(
         # sequential order
         if sequence_length > 1:
             if mini_batches > 1:
-                batches = np.array_split(self.all_sequence_indexes, len(self.all_sequence_indexes) // mini_batches)
+                batches = np.array_split(self.all_sequence_indexes, mini_batches)
                 return [[self.tensors_view[name][batch] for name in names] for batch in batches]
             return [[self.tensors_view[name][self.all_sequence_indexes] for name in names]]