Fix scoring with batches of size 1 by avoid squeezing the batch dim (#…

…1044)

CHANGELOG.md

@@ -11,6 +11,12 @@ Note that Sockeye has checks in place to not translate with an old model that wa
 
 Each version section may have subsections for: _Added_, _Changed_, _Removed_, _Deprecated_, and _Fixed_.
 
+## [3.1.12]
+
+### Fixed
+
+- Fix scoring with batches of size 1 (whic may occur when `|data| % batch_size == 1`.
+
 ## [3.1.11]
 
 ### Fixed

sockeye/__init__.py

@@ -11,4 +11,4 @@
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
-__version__ = '3.1.11'
+__version__ = '3.1.12'

sockeye/scoring.py

@@ -64,7 +64,7 @@ def forward(self,
 
         # Select the label log probability
         # logprobs and scores: (batch_size, target_seq_len)
-        token_scores = logprobs.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze()
+        token_scores = logprobs.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1)
         if self.score_type == C.SCORING_TYPE_NEGLOGPROB:
             token_scores = -token_scores
 
@@ -80,7 +80,7 @@ def forward(self,
             factor_scores = []  # type: List[pt.Tensor]
             for factor_logit, factor_label in factor_logits_and_labels:
                 factor_logprobs = factor_logit.log_softmax(dim=-1)
-                factor_token_scores = factor_logprobs.gather(dim=-1, index=factor_label.unsqueeze(-1)).squeeze()
+                factor_token_scores = factor_logprobs.gather(dim=-1, index=factor_label.unsqueeze(-1)).squeeze(-1)
                 if self.score_type == C.SCORING_TYPE_NEGLOGPROB:
                     factor_token_scores = -factor_token_scores
                 fs = factor_token_scores.masked_fill_(factor_label == C.PAD_ID, .0).sum(dim=-1, keepdims=True)  # type: ignore