Merge pull request octo-models#139 from rail-berkeley/dibya_patch_1

Fix big woopise
evelynmitchell · Dec 8, 2023 · 71658de · 71658de
2 parents 127c99a + 98d27d1
commit 71658de
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/orca/data/dataset.py b/orca/data/dataset.py
@@ -259,11 +259,6 @@ def move_language_instruction_to_tasks(traj):
 
         dataset = dataset.map(move_language_instruction_to_tasks, num_parallel_calls)
 
-    if train and subsample_length is not None:
-        dataset = dataset.map(
-            partial(_subsample, subsample_length=subsample_length), num_parallel_calls
-        )
-
     dataset = dataset.map(
         partial(
             _chunk_act_obs,
@@ -273,6 +268,10 @@ def move_language_instruction_to_tasks(traj):
         ),
         num_parallel_calls,
     )
+    if train and subsample_length is not None:
+        dataset = dataset.map(
+            partial(_subsample, subsample_length=subsample_length), num_parallel_calls
+        )
 
     return dataset