[QOL] throw assertion error if env was already done

ai4co · Nov 8, 2023 · c5cebc2 · c5cebc2
1 parent 7de5787
commit c5cebc2
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/rl4co/models/zoo/common/autoregressive/decoder.py b/rl4co/models/zoo/common/autoregressive/decoder.py
@@ -192,6 +192,7 @@ def forward(
             outputs.append(log_p)
             actions.append(action)
 
+        assert len(outputs) > 0, "No outputs were collected because all environments were done. Check your initial state"
         outputs, actions = torch.stack(outputs, 1), torch.stack(actions, 1)
         if calc_reward:
             td.set("reward", env.get_reward(td, actions))