Add drop_last option to method fit of PyTorchClassifier

Signed-off-by: Beat Buesser <[email protected]>

Author: Beat Buesser

art/estimators/classification/pytorch.py

@@ -362,6 +362,7 @@ def fit(  # pylint: disable=W0221
         batch_size: int = 128,
         nb_epochs: int = 10,
         training_mode: bool = True,
+        drop_last: bool = False,
         **kwargs,
     ) -> None:
         """
@@ -373,6 +374,9 @@ def fit(  # pylint: disable=W0221
         :param batch_size: Size of batches.
         :param nb_epochs: Number of epochs to use for training.
         :param training_mode: `True` for model set to training mode and `'False` for model set to evaluation mode.
+        :param drop_last: Set to ``True`` to drop the last incomplete batch, if the dataset size is not divisible by the batch
+                   size. If ``False`` and the size of dataset is not divisible by the batch size, then the last batch
+                   will be smaller. (default: ``False``)
         :param kwargs: Dictionary of framework-specific arguments. This parameter is not currently supported for PyTorch
                and providing it takes no effect.
         """
@@ -392,7 +396,11 @@ def fit(  # pylint: disable=W0221
         # Check label shape
         y_preprocessed = self.reduce_labels(y_preprocessed)
 
-        num_batch = int(np.floor(len(x_preprocessed) / float(batch_size)))
+        num_batch = len(x_preprocessed) / float(batch_size)
+        if drop_last:
+            num_batch = int(np.floor(num_batch))
+        else:
+            num_batch = int(np.ceil(num_batch))
         ind = np.arange(len(x_preprocessed))
 
         x_preprocessed = torch.from_numpy(x_preprocessed).to(self._device)