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 13, 2018
1 parent 527de38 commit 3cd1d01
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 @@ -1021,7 +1021,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 3cd1d01

Please sign in to comment.