From 16a99e770385455a38e966d8d862afb304ddc4ab Mon Sep 17 00:00:00 2001 From: Can Balioglu Date: Fri, 10 Jan 2025 00:01:06 +0000 Subject: [PATCH] Fix config_kls bug in SamplerHandlers --- src/fairseq2/generation/sampling/sampler.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/fairseq2/generation/sampling/sampler.py b/src/fairseq2/generation/sampling/sampler.py index 2b0b26d40..f5db6bff5 100644 --- 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