tree_map() missing 1 required positional argument: 'tree'

[lishulincug]

TypeError Traceback (most recent call last)
in
22 train_eval_iter = datasource.create_iterator(
23 shuffled(datasource.train_ids), batch_size=0, devices=devices)
---> 24 val_eval_iter = datasource.create_iterator(
25 shuffled(datasource.val_ids), batch_size=0, devices=devices)

/usr/local/lib/python3.8/site-packages/hypernerf/datasets/core.py in create_iterator(self, item_ids, batch_size, repeat, flatten, shuffle, prefetch_size, shuffle_buffer_size, devices)
386 devices: Optional[Sequence[Any]] = None):
387 """Creates a dataset iterator for JAX."""
--> 388 dataset = self.create_dataset(
389 item_ids,
390 flatten=flatten,

/usr/local/lib/python3.8/site-packages/hypernerf/datasets/core.py in create_dataset(self, item_ids, flatten, shuffle, row_shuffle_buffer_size, pixel_shuffle_buffer_size)
406 """Creates a tf.data.Dataset."""
407 if self.preload:
--> 408 return self._create_preloaded_dataset(
409 item_ids, flatten=flatten, shuffle=shuffle)
410 else:

/usr/local/lib/python3.8/site-packages/hypernerf/datasets/core.py in _create_preloaded_dataset(self, item_ids, flatten, shuffle)
434 data_list = utils.parallel_map(load_fn, item_ids)
435 data_list = [_camera_to_rays_fn(item) for item in data_list]
...
--> 340 return tree_util.tree_multimap(lambda *x: np.stack(x), *list_of_pytrees)
341
342

TypeError: tree_map() missing 1 required positional argument: 'tree'