From c9baae229898cf117f96dc9e61f04024ab9a9183 Mon Sep 17 00:00:00 2001 From: Amer Elsheikh Date: Tue, 29 Sep 2026 02:14:10 -0700 Subject: [PATCH] Automated Code Change PiperOrigin-RevId: 990158304 --- optformer/common/data/datasets/generator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optformer/common/data/datasets/generator.py b/optformer/common/data/datasets/generator.py index 09bb3d7..8c81b03 100644 --- a/optformer/common/data/datasets/generator.py +++ b/optformer/common/data/datasets/generator.py @@ -30,7 +30,7 @@ class GeneratorDatasetFn(base.DatasetFn[_S]): featurizer: featurizers.Featurizer[_S] = attrs.field(init=True, kw_only=True) - def __call__(self, source: Iterator[_S]) -> tf.data.Dataset: + def __call__(self, source: Iterator[_S]) -> tf.data.Dataset: # pyrefly: ignore[bad-override] def _generator(): for obj in source: yield self.featurizer.to_features(obj)