Fix sum_independent_dims docstring to reflect output shape (#1851)

Co-authored-by: Heinrick Lumini <[email protected]>
DLR-RM · Feb 27, 2024 · 56f20e4 · 56f20e4
1 parent a8e9059
commit 56f20e4
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst
@@ -65,6 +65,7 @@ Documentation:
 - Updated export to ONNX documentation, it is now much simpler to export SB3 models with newer ONNX Opset!
 - Added video link to "Practical Tips for Reliable Reinforcement Learning" video
 - Added ``render_mode="human"`` in the README example (@marekm4)
+- Fixed docstring signature for sum_independent_dims (@stagoverflow)
 
 Release 2.2.1 (2023-11-17)
 --------------------------
@@ -1565,4 +1566,4 @@ And all the contributors:
 @anand-bala @hughperkins @sidney-tio @AlexPasqua @dominicgkerr @Akhilez @Rocamonde @tobirohrer @ZikangXiong @ReHoss
 @DavyMorgan @luizapozzobon @Bonifatius94 @theSquaredError @harveybellini @DavyMorgan @FieteO @jonasreiher @npit @WeberSamuel @troiganto
 @lutogniew @lbergmann1 @lukashass @BertrandDecoster @pseudo-rnd-thoughts @stefanbschneider @kyle-he @PatrickHelm @corentinlger
-@marekm4
+@marekm4 @stagoverflow
diff --git a/stable_baselines3/common/distributions.py b/stable_baselines3/common/distributions.py
@@ -113,7 +113,7 @@ def sum_independent_dims(tensor: th.Tensor) -> th.Tensor:
     so we can sum components of the ``log_prob`` or the entropy.
 
     :param tensor: shape: (n_batch, n_actions) or (n_batch,)
-    :return: shape: (n_batch,)
+    :return: shape: (n_batch,) for (n_batch, n_actions) input, scalar for (n_batch,) input
     """
     if len(tensor.shape) > 1:
         tensor = tensor.sum(dim=1)