use type specified in args for softlogmax function

bd6718e4 · Nathan Habib · bff97e59 · bd6718e4
Commit bd6718e4 authored Jul 03, 2024 by Nathan Habib
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

lm_eval/models/huggingface.py lm_eval/models/huggingface.py +1 -1

No files found.
--- a/lm_eval/models/huggingface.py
+++ b/lm_eval/models/huggingface.py
@@ -1128,7 +1128,7 @@ class HFLM(TemplateLM):
            multi_logits = F.log_softmax(
                logits,
                dim=-1,
-                # dtype=logits.dtype,
+                dtype=logits.dtype,
            )  # [batch, padding_length (inp or cont), vocab]

            for (request_str, ctx_tokens, _), logits, inplen, cont_toks in zip(