Setting mixture_or_task variable correctly.

PiperOrigin-RevId: 706690426

protoscribe/pmmx/utils/seqio_utils.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """General utility functions for t5x."""
+
 import time
 from typing import Any, Callable, Optional, Type
 
@@ -87,6 +88,11 @@ def get_dataset(
         cfg.mixture_or_task_name not in mixtures):
       define_task_fn(cfg.mixture_or_task_name)
 
+  if isinstance(cfg.mixture_or_task_name, seqio.DatasetProviderBase):
+    mixture_or_task = cfg.mixture_or_task_name
+  else:
+    mixture_or_task = seqio.get_mixture_or_task(cfg.mixture_or_task_name)
+
   shard_info = seqio.ShardInfo(index=shard_id, num_shards=num_shards)
 
   if cfg.seed is None: