Skip to content
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions lmdeploy/pytorch/engine/model_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -401,6 +401,7 @@ def warmup(self):
is_decoding=False,
device='cuda',
vocab_size=self.model_config.vocab_size)
inputs.build_dp_meta()
self._forward_impl(inputs)

# warmup decoding(with cuda graph)
Expand All @@ -411,6 +412,7 @@ def warmup(self):
is_decoding=True,
device='cuda',
vocab_size=self.model_config.vocab_size)
inputs.build_dp_meta()
self._forward_impl(inputs)

def _slice_outs(self, inputs: torch.Tensor, seq_length: torch.LongTensor):
Expand Down