remove unused code

octo/data/traj_transforms.py

@@ -108,18 +108,6 @@ def subsample(traj: dict, subsample_length: int) -> dict:
     return traj
 
 
-def zero_out_future_proprio(traj: dict) -> dict:
-    """Removes all proprio inputs after first one to prevent causal confusion."""
-    traj["observation"]["proprio"] = tf.concat(
-        (
-            traj["observation"]["proprio"][:, :1],
-            tf.zeros_like(traj["observation"]["proprio"][:, 1:]),
-        ),
-        axis=1,
-    )
-    return traj
-
-
 def add_pad_mask_dict(traj: dict) -> dict:
     """Adds a dictionary indicating which elements of the observation/task should be treated as padding.
 

scripts/configs/finetune_config.py

@@ -44,8 +44,6 @@ def get_config(config_string="full,multimodal"):
             "heads_*.map_head.probe",
             "heads_*.map_head.MultiHeadDotProductAttention_0.*",
         )
-    elif mode == "frozen_transformer":
-        frozen_keys = ("octo_transformer.BlockTransformer_0.*",)
     else:
         raise ValueError("Invalid mode")
 

scripts/train.py

@@ -58,7 +58,9 @@ def main(_):
     jax_utils.initialize_compilation_cache()
 
     assert FLAGS.config.dataset_kwargs.batch_size % jax.device_count() == 0
+    assert FLAGS.config.viz_kwargs.eval_batch_size % jax.device_count() == 0
     assert FLAGS.config.dataset_kwargs.batch_size % jax.process_count() == 0
+    assert FLAGS.config.viz_kwargs.eval_batch_size % jax.process_count() == 0
 
     # create a 1D mesh with a single axis named "batch"
     mesh = Mesh(jax.devices(), axis_names="batch")