sampling should work now

r-pad · Aug 9, 2023 · 9113a69 · 9113a69
1 parent db00a0e
commit 9113a69
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/flowbot3d/datasets/flow_dataset_pyg.py b/flowbot3d/datasets/flow_dataset_pyg.py
@@ -48,8 +48,10 @@ def get_processed_dir(randomize_joints, randomize_camera):
  camera_chunk = "rc" if randomize_camera else "sc"
  return f"processed_{joint_chunk}_{camera_chunk}"
 
- def get_data(self, obj_id: str) -> Flowbot3DTGData:
- data_dict = self.dataset.get_data(obj_id, self.seed)
+ def get_data(self, obj_id: str, seed=None) -> Flowbot3DTGData:
+ if seed is None:
+ seed = self.seed
+ data_dict = self.dataset.get_data(obj_id, seed)
 
  data = tgd.Data(
  id=data_dict["id"],