Commit bff97e59 authored by Nathan Habib's avatar Nathan Habib
Browse files

use type specified in args for softlogmax function

parent 775d5d85
...@@ -1124,10 +1124,11 @@ class HFLM(TemplateLM): ...@@ -1124,10 +1124,11 @@ class HFLM(TemplateLM):
"labels": batched_conts, "labels": batched_conts,
} }
logits = self._model_call(batched_inps, **call_kwargs)
multi_logits = F.log_softmax( multi_logits = F.log_softmax(
self._model_call(batched_inps, **call_kwargs), logits,
dim=-1, dim=-1,
# dtype=torch.float16, # dtype=logits.dtype,
) # [batch, padding_length (inp or cont), vocab] ) # [batch, padding_length (inp or cont), vocab]
for (request_str, ctx_tokens, _), logits, inplen, cont_toks in zip( for (request_str, ctx_tokens, _), logits, inplen, cont_toks in zip(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment