Fix squeeze inside NeuralNetBinaryClassifier.infer

CHANGES.md

@@ -31,6 +31,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Fixed
 
 - Fixed a bug that caused LoadInitState not to work with TrainEndCheckpoint (#528)
+- Fixed NeuralNetBinaryClassifier wrongly squeezing the batch dimension when using batch_size = 1 (#558)
 
 
 ## [0.6.0] - 2019-07-19

skorch/classifier.py

@@ -307,7 +307,7 @@ def infer(self, x, **fit_params):
 
         The first output of the ``module`` must be a single array that
         has either shape (n,) or shape (n, 1). In the latter case, the
-        output will be squeezed to become 1-dim.
+        output will be reshaped to become 1-dim.
 
         """
         y_infer = super().infer(x, **fit_params)
@@ -320,7 +320,7 @@ def infer(self, x, **fit_params):
                 "Expected module output to have shape (n,) or "
                 "(n, 1), got {} instead".format(tuple(y_infer.shape)))
 
-        y_infer = y_infer.squeeze()
+        y_infer = y_infer.reshape(-1)
         if rest is None:
             return y_infer
         return (y_infer,) + tuple(rest)

skorch/tests/test_classifier.py

@@ -235,6 +235,15 @@ def test_net_learns(self, net_cls, module_cls, data):
         valid_acc = net.history[-1, 'valid_acc']
         assert valid_acc > 0.65
 
+    def test_batch_size_one(self, net_cls, module_cls, data):
+        X, y = data
+        net = net_cls(
+            module_cls,
+            max_epochs=1,
+            batch_size=1,
+        )
+        net.fit(X, y)
+
     def test_history_default_keys(self, net_fit):
         expected_keys = {
             'train_loss', 'valid_loss', 'epoch', 'dur', 'batches', 'valid_acc'