Fix factor map op for shortlist

rhenry-nv · Jul 9, 2021 · 7be492a · 7be492a
1 parent 22d13b3
commit 7be492a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/src/graph/node_operators_binary.h b/src/graph/node_operators_binary.h
@@ -1291,9 +1291,9 @@ size_t numLemmas_;
 bool hasShortlist_;
 public:
   AddFactorMaxesOp(const std::vector<Expr>& nodes, bool hasShortlist, size_t groupStart, size_t numLemmas)
-      : NaryNodeOp(nodes, getShape(nodes, hasShortlist), commonType(std::vector<Expr>(nodes.begin() + 1, nodes.end())) ) {
+      : NaryNodeOp(nodes, getShape(nodes, hasShortlist), commonType(std::vector<Expr>(nodes.begin() + 1 + (int)hasShortlist, nodes.end())) ) {
     groupStart_ = groupStart;
-    numLemmas_ = hasShortlist? nodes[1]->shape().size(): numLemmas;
+    numLemmas_ = numLemmas;
     hasShortlist_ = hasShortlist;
   }
 

diff --git a/src/layers/generic.cpp b/src/layers/generic.cpp
@@ -89,7 +89,7 @@ namespace marian {
     } else {
       auto numGroups = getNumFactorGroups();
       if(numGroups > 1 && graph()->isInference() && graph()->getBackend()->getDeviceId().type == DeviceType::gpu) {
-        Expr shortlistIndices = shortlist? constant(shortlist->indices()) : nullptr;
+        Expr shortlistIndices = shortlist? indices(shortlist->indices()) : nullptr;
         Expr lemmaHasFactorGroupTensor = getLemmaHasFactorGroupTensor();
         std::vector<Expr> groupLosses(logits_.size());
         std::transform(logits_.begin(), logits_.end(), groupLosses.begin(), [](const Ptr<RationalLoss>& loss) -> Expr {return loss->loss();});