Unverified Commit 98f97133 authored by Joe Runde's avatar Joe Runde Committed by GitHub
Browse files

[Bugfix] Fix TP > 1 for new granite (#8544)


Signed-off-by: default avatarJoe Runde <Joseph.Runde@ibm.com>
parent 56c3de01
......@@ -428,6 +428,7 @@ class GraniteForCausalLM(nn.Module, SupportsLoRA):
sampling_metadata: SamplingMetadata) -> Optional[torch.Tensor]:
logits = self.logits_processor(self.lm_head, hidden_states,
sampling_metadata)
if logits is not None:
logits /= self.config.logits_scaling
return logits
......
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