Skip to content

Commit

Permalink
[TextGeneration][Bug] Fix timer (#1294)
Browse files Browse the repository at this point in the history
* create new timer if none

* move timer check/creation to timer.py
  • Loading branch information
dsikka committed Oct 4, 2023
1 parent 93ba723 commit 5d022a7
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 1 deletion.
2 changes: 1 addition & 1 deletion src/deepsparse/transformers/engines/nl_decoder_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -185,7 +185,7 @@ def __call__(
:return: The generated token and corresponding logits
"""
timer = self.timer_manager.current
timer = self.timer_manager.current_or_new()
if kv_cache:
# if model has kv cache enabled, we need
# to add the kv cache state to the input
Expand Down
10 changes: 10 additions & 0 deletions src/deepsparse/utils/timer.py
Original file line number Diff line number Diff line change
Expand Up @@ -338,6 +338,16 @@ def all_times(self) -> Dict[str, List[float]]:

return all_times

def current_or_new(self) -> StagedTimer:
"""
Return the current timer if there is one, otherwise return a new one.
"""
if self.current:
return self.current
else:
with self.new_timer_context(total_inference=False) as timer:
return timer

def clear(self):
for t in self._timers:
t.clear()
Expand Down

0 comments on commit 5d022a7

Please sign in to comment.