amend

pytorch · Apr 18, 2024 · 5152b9d · 5152b9d
1 parent 7c7d4ce
commit 5152b9d
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/torchrl/objectives/value/functional.py b/torchrl/objectives/value/functional.py
@@ -1378,6 +1378,9 @@ def reward2go(
     cumsum = _inv_pad_sequence(cumsum, num_per_traj)
     cumsum = cumsum.reshape_as(reward)
     cumsum = cumsum.transpose(-2, -1)
-    # if cumsum.shape != shape:
+    if cumsum.shape != shape:
+        raise RuntimeError(
+            f"Wrong shape for output reward2go: {cumsum.shape} when {shape} was expected."
+        )
     #     cumsum = cumsum.view(shape)
     return cumsum