From e79a8d1adb31b2cf03113f14f329166de3d8c83b Mon Sep 17 00:00:00 2001
From: Ben Eisner <ben.a.eisner@gmail.com>
Date: Tue, 16 Apr 2024 14:07:05 -0400
Subject: [PATCH] Update dataset.py

---
 src/rpad/partnet_mobility_utils/dataset.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/rpad/partnet_mobility_utils/dataset.py b/src/rpad/partnet_mobility_utils/dataset.py
index 1539ef7..ffc56d6 100644
--- 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