From 1a9da6d7131fb9c90fc88382d05ce3c54834dbed Mon Sep 17 00:00:00 2001
From: divyashreepathihalli <divyashreepathihalli@gmail.com>
Date: Thu, 5 Mar 2026 06:11:41 +0000
Subject: [PATCH] fix tfds bug

---
 acme/datasets/tfds.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/acme/datasets/tfds.py b/acme/datasets/tfds.py
index 5c6411ec40..cecdc66f7c 100644
--- a/acme/datasets/tfds.py
+++ b/acme/datasets/tfds.py
@@ -200,10 +200,9 @@ def sample_and_postprocess(key: jnp.ndarray) -> Tuple[Any, jnp.ndarray]:
         data, key = pmapped_sample(self._jax_dataset, key)
         # All pmapped devices return the same data, so we just take the one from
         # the first device.
+        from acme.jax import utils as jax_utils_acme
         return (
-            jax.tree_util.tree_map(
-                lambda x: x.addressable_shards[0].data.squeeze(0), data
-            ),
+            jax_utils_acme.get_from_first_device(data, as_numpy=False),
             key,
         )
       self._sample = sample_and_postprocess