add pickable

tests/base/test_datasets.py

@@ -0,0 +1,19 @@
+import pickle
+
+import cloudpickle
+import pytest
+
+from tests.base.datasets import MNIST, TrialMNIST, AverageDataset
+
+
+@pytest.mark.parametrize('dataset_cls', [MNIST, TrialMNIST, AverageDataset])
+def test_pickling_dataset_mnist(tmpdir, dataset_cls):
+    mnist = dataset_cls()
+
+    mnist_pickled = pickle.dumps(mnist)
+    mnist_loaded = pickle.loads(mnist_pickled)
+    # assert vars(mnist) == vars(mnist_loaded)
+
+    mnist_pickled = cloudpickle.dumps(mnist)
+    mnist_loaded = cloudpickle.loads(mnist_pickled)
+    # assert vars(mnist) == vars(mnist_loaded)

tests/callbacks/test_callbacks.py

@@ -1,3 +1,4 @@
+import pickle
 from pathlib import Path
 
 import pytest
@@ -291,7 +292,6 @@ def training_step(self, *args, **kwargs):
 
 
 def test_pickling(tmpdir):
-    import pickle
     early_stopping = EarlyStopping()
     ckpt = ModelCheckpoint(tmpdir)