Make FBetaMeasure work with batch size of 1. (#2777)

- A stray squeeze caused batches of size one to be zero dimensional. Indexing into them was thus invalid and crashed.
- See #2275 (comment)

Author: brendan-ai2

allennlp/tests/training/metrics/fbeta_measure_test.py

@@ -161,3 +161,17 @@ def test_fbeta_multiclass_with_explicit_labels(self):
         numpy.testing.assert_almost_equal(precisions, desired_precisions, decimal=2)
         numpy.testing.assert_almost_equal(recalls, desired_recalls, decimal=2)
         numpy.testing.assert_almost_equal(fscores, desired_fscores, decimal=2)
+
+    def test_fbeta_handles_batch_size_of_one(self):
+        predictions = torch.Tensor([[0.2862, 0.3479, 0.1627, 0.2033]])
+        targets = torch.Tensor([1])
+        mask = torch.Tensor([1])
+
+        fbeta = FBetaMeasure()
+        fbeta(predictions, targets, mask)
+        metric = fbeta.get_metric()
+        precisions = metric['precision']
+        recalls = metric['recall']
+
+        numpy.testing.assert_almost_equal(precisions, [0.0, 1.0, 0.0, 0.0])
+        numpy.testing.assert_almost_equal(recalls, [0.0, 1.0, 0.0, 0.0])

allennlp/training/metrics/fbeta_measure.py

@@ -122,7 +122,7 @@ def __call__(self,
         mask = mask.to(torch.uint8)
         gold_labels = gold_labels.float()
 
-        argmax_predictions = predictions.max(dim=-1)[1].float().squeeze(dim=-1)
+        argmax_predictions = predictions.max(dim=-1)[1].float()
         true_positives = (gold_labels == argmax_predictions) * mask
         true_positives_bins = gold_labels[true_positives]