[BugFix] Fix RLHF (pytorch#1757)

ai4co · Dec 23, 2023 · 15950d1 · 15950d1
1 parent 64434df
commit 15950d1
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/test/assets/tldr_batch.zip b/test/assets/tldr_batch.zip
diff --git a/test/test_rlhf.py b/test/test_rlhf.py
@@ -66,6 +66,9 @@ def tldr_batch_dir(tmp_path_factory):
     with zipfile.ZipFile(dataset_path, "r") as zip_ref:
         zip_ref.extractall(dest)
         yield dest / Path(dataset_path).stem
+    from torchrl._utils import print_directory_tree
+
+    print_directory_tree(dest)
 
 
 @pytest.mark.skipif(
@@ -436,7 +439,7 @@ def init_reward_model(device=None):
 
     @staticmethod
     def _get_dummy_batch(batch_dir):
-        return PromptData.from_tensordict(TensorDict.load_memmap(batch_dir))
+        return TensorDict.load_memmap(batch_dir)
 
     @property
     def _model(self):