making get_screenshots work for png and jpg

ServiceNow · Oct 4, 2024 · 49ada25 · 49ada25
1 parent c42496b
commit 49ada25
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/browsergym/experiments/src/browsergym/experiments/loop.py b/browsergym/experiments/src/browsergym/experiments/loop.py
@@ -614,7 +614,7 @@ def get_screenshot(self, step: int, som=False) -> Image:
         return self._screenshots[key]
 
     def get_screenshots(self, som=False):
-        files = list(self.exp_dir.glob("screenshot_step_*.png"))
+        files = list(self.exp_dir.glob("screenshot_step_*"))
         max_step = 0
         for file in files:
             step = int(file.name.split("_")[-1].split(".")[0])