Skip to content

Commit

Permalink
Change to use MXNet's topk for CPUs in inference
Browse files Browse the repository at this point in the history
Since MXNet's topk has better performance than numpy version with
PR apache/mxnet#12085, in order
to leverage such performance boost, change to use MXNet's topk for
CPU device when doing inference.
  • Loading branch information
ciyongch committed Aug 23, 2018
1 parent 241a6c8 commit cac13da
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion sockeye/inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -1009,7 +1009,7 @@ def __init__(self,
k=self.beam_size,
batch_size=self.batch_size,
offset=self.offset,
use_mxnet_topk=self.context != mx.cpu()) # MXNet implementation is faster on GPUs
use_mxnet_topk=True)

self._sort_by_index = SortByIndex()
self._sort_by_index.initialize(ctx=self.context)
Expand Down

0 comments on commit cac13da

Please sign in to comment.