feat(overcooked): more viz JIT

FLAIROx · Jun 29, 2024 · 5920c21 · 5920c21
1 parent 0182667
commit 5920c21
Show file tree

Hide file tree

Showing 2 changed files with 122 additions and 146 deletions.
diff --git a/jaxmarl/viz/grid_rendering.py b/jaxmarl/viz/grid_rendering.py
@@ -12,8 +12,8 @@ def downsample(img, factor):
     img = img.reshape(
         [img.shape[0] // factor, factor, img.shape[1] // factor, factor, 3]
     )
-    img = img.mean(axis=3)
-    img = img.mean(axis=1)
+    img = img.mean(axis=3, dtype=jnp.uint8)
+    img = img.mean(axis=1, dtype=jnp.uint8)
 
     return img