Merge pull request #77 from neuro-ml/develop

Batch iter compatibility with lightning

dpipe/__version__.py

@@ -1 +1 @@
-__version__ = '0.2.4'
+__version__ = '0.2.5'

dpipe/batch_iter/pipeline.py

@@ -131,6 +131,25 @@ def close(self):
         """Stop all background processes."""
         self.__exit__(None, None, None)
 
+    @property
+    def closing_callback(self):
+        """
+        A callback to make this interface compatible with `Lightning` which allows for a safe release of resources
+
+        Examples
+        --------
+        >>> batch_iter = Infinite(...)
+        >>> trainer = Trainer(callbacks=[batch_iter.closing_callback, ...])
+        """
+        from lightning.pytorch.callbacks import Callback
+
+        class ClosingCallback(Callback):
+            def teardown(self, trainer, pl_module, stage):
+                this.close()
+
+        this = self
+        return ClosingCallback()
+
     def __iter__(self):
         return self()