Update dataset.py

r-pad · Apr 16, 2024 · e79a8d1 · e79a8d1
1 parent c29145f
commit e79a8d1
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/rpad/partnet_mobility_utils/dataset.py b/src/rpad/partnet_mobility_utils/dataset.py
@@ -56,6 +56,7 @@ def __init__(
         root: Union[str, Path],
         split: Union[AVAILABLE_DATASET, List[str]],
         renderer: Literal["pybullet", "sapien", "trimesh"] = "pybullet",
+        use_egl: bool = False,
     ):
         if isinstance(split, str):
             if split == "all":
@@ -88,6 +89,7 @@ def __init__(
         }
         self.renderers: Dict[str, PMRenderer] = {}
         self.renderer_type = renderer
+        self.use_egl = use_egl
 
     def get(
         self,
@@ -117,6 +119,7 @@ def get(
             joints=joints,
             camera_xyz=camera_xyz,
             seed=seed,
+            use_egl=use_egl,
         )
 
         return pc_render