Fix stat tracking bug

2023-06-26 22:25:43 -07:00
parent 5c16a90ceb
commit 1ce0994c8a
1 changed files with 2 additions and 0 deletions
--- a/ddpo_pytorch/stat_tracking.py
+++ b/ddpo_pytorch/stat_tracking.py
@@ -9,6 +9,8 @@ class PerPromptStatTracker:
        self.stats = {}

    def update(self, prompts, rewards):
+        prompts = np.array(prompts)
+        rewards = np.array(rewards)
        unique = np.unique(prompts)
        advantages = np.empty_like(rewards)
        for prompt in unique: