Unwrap accelerated model on train end (#842)

It is recommended to "unwrap" the model after training ends:

https://github.com/huggingface/accelerate/blob/6ffab178accebda485295bddf8eb6bf436ff698f/src/accelerate/accelerator.py#L542-L551

https://github.com/huggingface/accelerate/blob/e549cea65ced232131d2d028717b6a0057274add/docs/source/quicktour.rst#savingloading-a-model

skorch/helper.py

@@ -669,3 +669,8 @@ def _step_optimizer(self, step_fn):
         for name in self._optimizers:
             optimizer = getattr(self, name + '_')
             optimizer.step()
+
+    # pylint: disable=unused-argument
+    def on_train_end(self, net, X=None, y=None, **kwargs):
+        super().on_train_end(net, X=X, y=y, **kwargs)
+        self.module_ = self.accelerator.unwrap_model(self.module_)

skorch/tests/test_helper.py

@@ -835,6 +835,9 @@ def backward(self, loss, **kwargs):
                 loss.backward(**kwargs)
                 loss.backward_was_called = True
 
+            def unwrap_model(self, model):
+                return model
+
         # pylint: disable=missing-class-docstring
         class AcceleratedNet(AccelerateMixin, NeuralNetClassifier):
             def get_iterator(self, *args, **kwargs):