Fix config_kls bug in SamplerHandlers

facebookresearch · Jan 10, 2025 · 16a99e7 · 16a99e7
1 parent 6bfc6d0
commit 16a99e7
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/fairseq2/generation/sampling/sampler.py b/src/fairseq2/generation/sampling/sampler.py
@@ -150,7 +150,7 @@ def create(self, config: object) -> Sampler:
     @property
     @override
     def config_kls(self) -> type:
-        return TopPSamplerHandler
+        return TopPSamplerConfig
 
 
 TOP_K_SAMPLER: Final = "top-k"
@@ -172,4 +172,4 @@ def create(self, config: object) -> Sampler:
     @property
     @override
     def config_kls(self) -> type:
-        return TopKSamplerHandler
+        return TopKSamplerConfig