diff --git a/.github/workflows/deepinterpolation.yml b/.github/workflows/deepinterpolation.yml
new file mode 100644
index 0000000000..7f5e3e54f9
--- /dev/null
+++ b/.github/workflows/deepinterpolation.yml
@@ -0,0 +1,51 @@
+name: Testing deepinterpolation
+
+on:
+  pull_request:
+    types: [synchronize, opened, reopened]
+    branches:
+      - main
+
+concurrency:  # Cancel previous workflows on the same pull request
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  build-and-test:
+    name: Test on ${{ matrix.os }} OS
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: ["ubuntu-latest"]
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        with:
+          python-version: '3.9'
+      - name: Get changed files
+        id: changed-files
+        uses: tj-actions/changed-files@v35
+      - name: Deepinteprolation changes
+        id: modules-changed
+        run: |
+          for file in ${{ steps.changed-files.outputs.all_changed_files }}; do
+            if [[ $file == *"/deepinterpolation/"* ]]; then
+              echo "DeepInterpolation changed"
+              echo "DEEPINTERPOLATION_CHANGED=true" >> $GITHUB_OUTPUT
+            fi
+          done
+      - name: Install dependencies
+        if: ${{ steps.modules-changed.outputs.DEEPINTERPOLATION_CHANGED == 'true' }}
+        run: |
+          python -m pip install -U pip  # Official recommended way
+          # install deepinteprolation
+          pip install tensorflow==2.7.0
+          pip install deepinterpolation@git+https://github.com/AllenInstitute/deepinterpolation.git
+          pip install protobuf==3.20.*
+          pip install -e .[full,test_core]
+      - name: Test DeepInterpolation with pytest
+        if: ${{ steps.modules-changed.outputs.DEEPINTERPOLATION_CHANGED == 'true' }}
+        run: |
+          pytest -v src/spikeinterface/preprocessing/deepinterpolation
+        shell: bash  # Necessary for pipeline to work on windows
diff --git a/.github/workflows/full-test.yml b/.github/workflows/full-test.yml
index 8f88e84039..dad42e021b 100644
--- a/.github/workflows/full-test.yml
+++ b/.github/workflows/full-test.yml
@@ -136,7 +136,7 @@ jobs:
         run: ./.github/run_tests.sh "extractors and not streaming_extractors"
       - name: Test preprocessing
         if: ${{ steps.modules-changed.outputs.PREPROCESSING_CHANGED == 'true' || steps.modules-changed.outputs.CORE_CHANGED == 'true' }}
-        run: ./.github/run_tests.sh preprocessing
+        run: ./.github/run_tests.sh "preprocessing and not deepinterpolation"
       - name: Test postprocessing
         if: ${{ steps.modules-changed.outputs.POSTPROCESSING_CHANGED == 'true' || steps.modules-changed.outputs.CORE_CHANGED == 'true' }}
         run: ./.github/run_tests.sh postprocessing
diff --git a/src/spikeinterface/preprocessing/deepinterpolation/__init__.py b/src/spikeinterface/preprocessing/deepinterpolation/__init__.py
index 2ff166a81e..13d9c69def 100644
--- a/src/spikeinterface/preprocessing/deepinterpolation/__init__.py
+++ b/src/spikeinterface/preprocessing/deepinterpolation/__init__.py
@@ -1 +1,2 @@
 from .deepinterpolation import DeepInterpolatedRecording, deepinterpolate
+from .train import train_deepinterpolation
diff --git a/src/spikeinterface/preprocessing/deepinterpolation/deepinterpolation.py b/src/spikeinterface/preprocessing/deepinterpolation/deepinterpolation.py
index f82ea8a5df..948ac9f224 100644
--- a/src/spikeinterface/preprocessing/deepinterpolation/deepinterpolation.py
+++ b/src/spikeinterface/preprocessing/deepinterpolation/deepinterpolation.py
@@ -1,6 +1,7 @@
 import numpy as np
-import os
+from typing import Optional
 
+from .tf_utils import has_tf, import_tf
 from ...core import BaseRecording
 from ...core.core_tools import define_function_from_class
 from ..basepreprocessor import BasePreprocessor, BasePreprocessorSegment
@@ -8,305 +9,111 @@
 from spikeinterface.core import get_random_data_chunks
 
 
-def import_tf(use_gpu=True, disable_tf_logger=True):
-    import tensorflow as tf
-
-    if not use_gpu:
-        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
-
-    if disable_tf_logger:
-        os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
-        tf.get_logger().setLevel("ERROR")
-
-    tf.compat.v1.disable_eager_execution()
-    gpus = tf.config.list_physical_devices("GPU")
-    if gpus:
-        try:
-            # Currently, memory growth needs to be the same across GPUs
-            for gpu in gpus:
-                tf.config.experimental.set_memory_growth(gpu, True)
-        except RuntimeError as e:
-            # Memory growth must be set before GPUs have been initialized
-            print(e)
-    return tf
-
-
-def has_tf(use_gpu=True, disable_tf_logger=True):
-    try:
-        import_tf(use_gpu, disable_tf_logger)
-        return True
-    except ImportError:
-        return False
+class DeepInterpolatedRecording(BasePreprocessor):
+    """
+    DeepInterpolatedRecording is a wrapper around a recording extractor that allows to apply a deepinterpolation model.
 
+    For more information, see:
+    Lecoq et al. (2021) Removing independent noise in systems neuroscience data using DeepInterpolation.
+    Nature Methods. 18: 1401-1408. doi: 10.1038/s41592-021-01285-2.
 
-def define_input_generator_class(use_gpu, disable_tf_logger=True):
-    """Define DeepInterpolationInputGenerator class at run-time
+    Parts of this code is adapted from https://github.com/AllenInstitute/deepinterpolation.
 
     Parameters
     ----------
-    use_gpu : bool
-        Whether to load TF with GPU capabilities
-    disable_tf_logger : bool, optional
-            If True, tensorflow logging is disabled, by default True
+    recording: BaseRecording
+        The recording extractor to be deepinteprolated
+    model_path: str
+        Path to the deepinterpolation h5 model
+    pre_frame: int
+        Number of frames before the frame to be predicted
+    post_frame: int
+        Number of frames after the frame to be predicted
+    pre_post_omission: int
+        Number of frames to be omitted before and after the frame to be predicted
+    batch_size: int
+        Batch size to be used for the prediction
+    predict_workers: int
+        Number of workers to be used for the tensorflow `predict` function.
+        Multiple workers can be used to speed up the prediction by pre-fetching the data.
+    use_gpu: bool
+        If True, the gpu will be used for the prediction
+    disable_tf_logger: bool
+        If True, the tensorflow logger will be disabled
+    memory_gpu: int
+        The amount of memory to be used by the gpu
 
     Returns
     -------
-    class
-        The defined DeepInterpolationInputGenerator class
+    recording: DeepInterpolatedRecording
+        The deepinterpolated recording extractor object
     """
-    tf = import_tf(use_gpu, disable_tf_logger)
-
-    class DeepInterpolationInputGenerator(tf.keras.utils.Sequence):
-        def __init__(
-            self,
-            recording,
-            start_frame,
-            end_frame,
-            batch_size,
-            pre_frames,
-            post_frames,
-            pre_post_omission,
-            local_mean,
-            local_std,
-        ):
-            self.recording = recording
-            self.start_frame = start_frame
-            self.end_frame = end_frame
-
-            self.batch_size = batch_size
-            self.last_batch_size = (end_frame - start_frame) - (self.__len__() - 1) * batch_size
-
-            self.pre_frames = pre_frames
-            self.post_frames = post_frames
-            self.pre_post_omission = pre_post_omission
-
-            self.local_mean = local_mean
-            self.local_std = local_std
-
-        def __len__(self):
-            return -((self.end_frame - self.start_frame) // -self.batch_size)
-
-        def __getitem__(self, idx):
-            n_batches = self.__len__()
-            if idx < n_batches - 1:
-                start_frame = self.start_frame + self.batch_size * idx - self.pre_frames - self.pre_post_omission
-                end_frame = self.start_frame + self.batch_size * (idx + 1) + self.post_frames + self.pre_post_omission
-                traces = self.recording.get_traces(
-                    start_frame=start_frame, end_frame=end_frame, channel_indices=slice(None)
-                )
-                batch_size = self.batch_size
-            else:
-                start_frame = self.end_frame - self.last_batch_size - self.pre_frames - self.pre_post_omission
-                end_frame = self.end_frame + self.post_frames + self.pre_post_omission
-                traces = self.recording.get_traces(
-                    start_frame=start_frame, end_frame=end_frame, channel_indices=slice(None)
-                )
-                batch_size = self.last_batch_size
-
-            shape = (traces.shape[0], int(384 / 2), 2)
-            traces = np.reshape(traces, newshape=shape)
-
-            di_input = np.zeros((batch_size, 384, 2, self.pre_frames + self.post_frames))
-            di_label = np.zeros((batch_size, 384, 2, 1))
-            for index_frame in range(
-                self.pre_frames + self.pre_post_omission, batch_size + self.pre_frames + self.pre_post_omission
-            ):
-                di_input[index_frame - self.pre_frames - self.pre_post_omission] = self.reshape_input_forward(
-                    index_frame, traces
-                )
-                di_label[index_frame - self.pre_frames - self.pre_post_omission] = self.reshape_label_forward(
-                    traces[index_frame]
-                )
-            return (di_input, di_label)
-
-        def reshape_input_forward(self, index_frame, raw_data):
-            """Reshapes the frames surrounding the target frame to the form expected by model;
-            also subtracts mean and divides by std.
-
-            Parameters
-            ----------
-            index_frame : int
-                index of the frame to be predicted
-            raw_data : ndarray; (frames, 192, 2)
-                a chunk of data used to generate the input
-
-            Returns
-            -------
-            input_full : ndarray; (1, 384, 2, pre_frames+post_frames)
-                input to trained network to predict the center frame
-            """
-            # currently only works for recordings with 384 channels
-            nb_probes = 384
-
-            # We reorganize to follow true geometry of probe for convolution
-            input_full = np.zeros([1, nb_probes, 2, self.pre_frames + self.post_frames], dtype="float32")
-
-            input_index = np.arange(
-                index_frame - self.pre_frames - self.pre_post_omission,
-                index_frame + self.post_frames + self.pre_post_omission + 1,
-            )
-            input_index = input_index[input_index != index_frame]
-
-            for index_padding in np.arange(self.pre_post_omission + 1):
-                input_index = input_index[input_index != index_frame - index_padding]
-                input_index = input_index[input_index != index_frame + index_padding]
-
-            data_img_input = raw_data[input_index, :, :]
-
-            data_img_input = np.swapaxes(data_img_input, 1, 2)
-            data_img_input = np.swapaxes(data_img_input, 0, 2)
-
-            even = np.arange(0, nb_probes, 2)
-            odd = even + 1
-
-            data_img_input = (data_img_input.astype("float32") - self.local_mean) / self.local_std
-
-            input_full[0, even, 0, :] = data_img_input[:, 0, :]
-            input_full[0, odd, 1, :] = data_img_input[:, 1, :]
-            return input_full
 
-        def reshape_label_forward(self, label):
-            """Reshapes the target frame to the form expected by model.
-
-            Parameters
-            ----------
-            label : ndarray, (1, 192, 2)
-
-            Returns
-            -------
-            reshaped_label : ndarray, (1, 384, 2, 1)
-                target frame after reshaping
-            """
-            # currently only works for recordings with 384 channels
-            nb_probes = 384
-
-            input_full = np.zeros([1, nb_probes, 2, 1], dtype="float32")
-
-            data_img_input = np.expand_dims(label, axis=0)
-            data_img_input = np.swapaxes(data_img_input, 1, 2)
-            data_img_input = np.swapaxes(data_img_input, 0, 2)
-
-            even = np.arange(0, nb_probes, 2)
-            odd = even + 1
-
-            data_img_input = (data_img_input.astype("float32") - self.local_mean) / self.local_std
-
-            input_full[0, even, 0, :] = data_img_input[:, 0, :]
-            input_full[0, odd, 1, :] = data_img_input[:, 1, :]
-            return input_full
-
-    return DeepInterpolationInputGenerator
-
-
-class DeepInterpolatedRecording(BasePreprocessor):
     name = "deepinterpolate"
 
     def __init__(
         self,
-        recording: BaseRecording,
+        recording,
         model_path: str,
-        pre_frames: int = 30,
-        post_frames: int = 30,
+        pre_frame: int = 30,
+        post_frame: int = 30,
         pre_post_omission: int = 1,
-        batch_size=128,
+        batch_size: int = 128,
         use_gpu: bool = True,
+        predict_workers: int = 1,
         disable_tf_logger: bool = True,
-        **random_chunk_kwargs,
+        memory_gpu: Optional[int] = None,
     ):
-        """Applies DeepInterpolation, a neural network based denoising method, to the recording.
-
-        Notes
-        -----
-        * Currently this only works on Neuropixels 1.0-like recordings with 384 channels.
-        If the recording has fewer number of channels, consider matching the channel count with
-        `ZeroChannelPaddedRecording`.
-        * The specified model must have the same input dimensions as the model from original paper.
-        * Will use GPU if available.
-        * Inference (application of model) is done lazily, i.e. only when `get_traces` is called.
-
-        For more information, see:
-        Lecoq et al. (2021) Removing independent noise in systems neuroscience data using DeepInterpolation.
-        Nature Methods. 18: 1401-1408. doi: 10.1038/s41592-021-01285-2.
-
-        Parts of this code is adapted from https://github.com/AllenInstitute/deepinterpolation.
+        assert has_tf(
+            use_gpu, disable_tf_logger, memory_gpu
+        ), "To use DeepInterpolation, you first need to install `tensorflow`."
 
-        Parameters
-        ----------
-        recording : si.BaseRecording
-        model_path : str
-            Path to pre-trained model
-        pre_frames : int
-            Number of frames before target frame used for training and inference
-        post_frames : int
-            Number of frames after target frame used for training and inference
-        pre_post_omission : int
-            Number of frames around the target frame to omit
-        batch_size : int, optional
-            Number of frames per batch to infer (adjust based on hardware); by default 128
-        disable_tf_logger : bool, optional
-            If True, tensorflow logging is disabled, by default True
-        random_chunk_kwargs: keyword arguments for get_random_data_chunks
-        """
-
-        assert has_tf(use_gpu, disable_tf_logger), "To use DeepInterpolation, you first need to install `tensorflow`."
-        assert recording.get_num_channels() <= 384, (
-            "DeepInterpolation only works on Neuropixels 1.0-like "
-            "recordings with 384 channels. This recording has too many "
-            "channels."
-        )
-        assert recording.get_num_channels() == 384, (
-            "DeepInterpolation only works on Neuropixels 1.0-like "
-            "recordings with 384 channels. "
-            "This recording has too few channels. Try matching the channel "
-            "count with `ZeroChannelPaddedRecording`."
-        )
-        self.tf = import_tf(use_gpu, disable_tf_logger)
+        self.tf = import_tf(use_gpu, disable_tf_logger, memory_gpu=memory_gpu)
 
         # try move model load here with spawn
         BasePreprocessor.__init__(self, recording)
 
         # first time retrieving traces check that dimensions are ok
         self.tf.keras.backend.clear_session()
-        self.model = self.tf.keras.models.load_model(filepath=model_path)
-        # check input shape for the last dimension
-        config = self.model.get_config()
-        input_shape = config["layers"][0]["config"]["batch_input_shape"]
-        assert input_shape[-1] == pre_frames + post_frames, (
-            "The sum of `pre_frames` and `post_frames` must match " "the last dimension of the model."
-        )
-
-        local_data = get_random_data_chunks(recording, **random_chunk_kwargs)
-        if isinstance(recording, ZeroChannelPaddedRecording):
-            local_data = local_data[:, recording.channel_mapping]
-
-        local_mean = np.mean(local_data.flatten())
-        local_std = np.std(local_data.flatten())
+        model = self.tf.keras.models.load_model(filepath=model_path)
+
+        # check shape (this will need to be done at inference)
+        network_input_shape = model.get_config()["layers"][0]["config"]["batch_input_shape"]
+        desired_shape = network_input_shape[1:3]
+        assert (
+            desired_shape[0] * desired_shape[1] == recording.get_num_channels()
+        ), "The desired shape of the network input must match the number of channels in the recording"
+        assert (
+            network_input_shape[-1] == pre_frame + post_frame
+        ), "The desired shape of the network input must match the pre and post frames"
 
+        self.model = model
         # add segment
         for segment in recording._recording_segments:
             recording_segment = DeepInterpolatedRecordingSegment(
                 segment,
                 self.model,
-                pre_frames,
-                post_frames,
+                pre_frame,
+                post_frame,
                 pre_post_omission,
-                local_mean,
-                local_std,
+                desired_shape,
                 batch_size,
-                use_gpu,
-                disable_tf_logger,
+                predict_workers,
             )
             self.add_recording_segment(recording_segment)
 
         self._preferred_mp_context = "spawn"
         self._kwargs = dict(
             recording=recording,
-            model_path=model_path,
-            pre_frames=pre_frames,
-            post_frames=post_frames,
+            model_path=str(model_path),
+            pre_frame=pre_frame,
+            post_frame=post_frame,
             pre_post_omission=pre_post_omission,
             batch_size=batch_size,
-            **random_chunk_kwargs,
+            predict_workers=predict_workers,
+            use_gpu=use_gpu,
+            disable_tf_logger=disable_tf_logger,
+            memory_gpu=memory_gpu,
         )
         self.extra_requirements.extend(["tensorflow"])
 
@@ -316,30 +123,26 @@ def __init__(
         self,
         recording_segment,
         model,
-        pre_frames,
-        post_frames,
+        pre_frame,
+        post_frame,
         pre_post_omission,
-        local_mean,
-        local_std,
+        desired_shape,
         batch_size,
-        use_gpu,
-        disable_tf_logger,
+        predict_workers,
     ):
         BasePreprocessorSegment.__init__(self, recording_segment)
 
         self.model = model
-        self.pre_frames = pre_frames
-        self.post_frames = post_frames
+        self.pre_frame = pre_frame
+        self.post_frame = post_frame
         self.pre_post_omission = pre_post_omission
-        self.local_mean = local_mean
-        self.local_std = local_std
         self.batch_size = batch_size
-        self.use_gpu = use_gpu
-
-        # creating class dynamically to use the imported TF with GPU enabled/disabled based on the use_gpu flag
-        self.DeepInterpolationInputGenerator = define_input_generator_class(use_gpu, disable_tf_logger)
+        self.desired_shape = desired_shape
+        self.predict_workers = predict_workers
 
     def get_traces(self, start_frame, end_frame, channel_indices):
+        from .generators import SpikeInterfaceRecordingSegmentGenerator
+
         n_frames = self.parent_recording_segment.get_num_samples()
 
         if start_frame == None:
@@ -350,16 +153,16 @@ def get_traces(self, start_frame, end_frame, channel_indices):
 
         # for frames that lack full training data (i.e. pre and post frames including omissinos),
         # just return uninterpolated
-        if start_frame < self.pre_frames + self.pre_post_omission:
-            true_start_frame = self.pre_frames + self.pre_post_omission
+        if start_frame < self.pre_frame + self.pre_post_omission:
+            true_start_frame = self.pre_frame + self.pre_post_omission
             array_to_append_front = self.parent_recording_segment.get_traces(
                 start_frame=0, end_frame=true_start_frame, channel_indices=channel_indices
             )
         else:
             true_start_frame = start_frame
 
-        if end_frame > n_frames - self.post_frames - self.pre_post_omission:
-            true_end_frame = n_frames - self.post_frames - self.pre_post_omission
+        if end_frame > n_frames - self.post_frame - self.pre_post_omission:
+            true_end_frame = n_frames - self.post_frame - self.pre_post_omission
             array_to_append_back = self.parent_recording_segment.get_traces(
                 start_frame=true_end_frame, end_frame=n_frames, channel_indices=channel_indices
             )
@@ -367,22 +170,24 @@ def get_traces(self, start_frame, end_frame, channel_indices):
             true_end_frame = end_frame
 
         # instantiate an input generator that can be passed directly to model.predict
-        input_generator = self.DeepInterpolationInputGenerator(
-            recording=self.parent_recording_segment,
+        batch_size = min(self.batch_size, true_end_frame - true_start_frame)
+        input_generator = SpikeInterfaceRecordingSegmentGenerator(
+            recording_segment=self.parent_recording_segment,
             start_frame=true_start_frame,
             end_frame=true_end_frame,
-            pre_frames=self.pre_frames,
-            post_frames=self.post_frames,
+            pre_frame=self.pre_frame,
+            post_frame=self.post_frame,
             pre_post_omission=self.pre_post_omission,
-            local_mean=self.local_mean,
-            local_std=self.local_std,
-            batch_size=self.batch_size,
+            batch_size=batch_size,
+            desired_shape=self.desired_shape,
         )
-        di_output = self.model.predict(input_generator, verbose=2)
+        di_output = self.model.predict(input_generator, workers=self.predict_workers, verbose=2)
 
-        out_traces = self.reshape_backward(di_output)
+        out_traces = input_generator.reshape_output(di_output)
 
-        if true_start_frame != start_frame:
+        if (
+            true_start_frame != start_frame
+        ):  # related to the restriction to be applied from the start and end frames around 0 and end
             out_traces = np.concatenate((array_to_append_front, out_traces), axis=0)
 
         if true_end_frame != end_frame:
@@ -390,31 +195,6 @@ def get_traces(self, start_frame, end_frame, channel_indices):
 
         return out_traces[:, channel_indices]
 
-    def reshape_backward(self, di_frames):
-        """reshapes the prediction from model back to frames
-
-        Parameters
-        ----------
-        di_frames : ndarray, (frames, 384, 2, 1)
-            predicted output of the model
-
-        Returns
-        -------
-        reshaped_frames : ndarray; (frames, 384)
-            predicted frames after reshaping
-        """
-        # currently works only for recording with 384 channels
-        nb_probes = 384
-        n_frames = di_frames.shape[0]
-        even = np.arange(0, nb_probes, 2)
-        odd = even + 1
-        reshaped_frames = np.zeros((n_frames, 384))
-        for frame in range(n_frames):
-            reshaped_frames[frame, 0::2] = di_frames[frame, even, 0, 0]
-            reshaped_frames[frame, 1::2] = di_frames[frame, odd, 1, 0]
-        reshaped_frames = reshaped_frames * self.local_std + self.local_mean
-        return reshaped_frames
-
 
 # function for API
 deepinterpolate = define_function_from_class(source_class=DeepInterpolatedRecording, name="deepinterpolate")
diff --git a/src/spikeinterface/preprocessing/deepinterpolation/generators.py b/src/spikeinterface/preprocessing/deepinterpolation/generators.py
new file mode 100644
index 0000000000..8200340ac1
--- /dev/null
+++ b/src/spikeinterface/preprocessing/deepinterpolation/generators.py
@@ -0,0 +1,311 @@
+from __future__ import annotations
+import tempfile
+import json
+from typing import Optional
+import numpy as np
+
+from ...core import load_extractor, concatenate_recordings, BaseRecording, BaseRecordingSegment
+
+from .tf_utils import has_tf, import_tf
+
+from ...core import load_extractor
+
+from deepinterpolation.generator_collection import SequentialGenerator
+
+
+class SpikeInterfaceRecordingGenerator(SequentialGenerator):
+    """
+    This generator is used when dealing with a SpikeInterface recording.
+    The desired shape controls the reshaping of the input data before convolutions.
+    """
+
+    def __init__(
+        self,
+        recordings: BaseRecording | list[BaseRecording],
+        pre_frame: int = 30,
+        post_frame: int = 30,
+        pre_post_omission: int = 1,
+        desired_shape: tuple = (192, 2),
+        batch_size: int = 100,
+        steps_per_epoch: int = 10,
+        start_frame: Optional[int] = None,
+        end_frame: Optional[int] = None,
+        total_samples: int = -1,
+    ):
+        if not isinstance(recordings, list):
+            recordings = [recordings]
+        self.recordings = recordings
+        if len(recordings) > 1:
+            assert (
+                r.get_num_channels() == recordings[0].get_num_channels() for r in recordings[1:]
+            ), "All recordings must have the same number of channels"
+
+        total_num_samples = np.sum(r.get_total_samples() for r in recordings)
+        # In case of multiple recordings and/or multiple segments, we calculate the frame periods to be excluded (borders)
+        exclude_intervals = []
+        pre_extended = pre_frame + pre_post_omission
+        post_extended = post_frame + pre_post_omission
+        for i, recording in enumerate(recordings):
+            total_samples_pre = np.sum(r.get_total_samples() for r in recordings[:i])
+            for segment_index in range(recording.get_num_segments()):
+                # exclude samples at the border of the recordings
+                num_samples_segment_pre = np.sum(recording.get_num_samples(s) for s in np.arange(segment_index))
+                if num_samples_segment_pre > 0:
+                    exclude_intervals.append(
+                        (
+                            total_samples_pre + num_samples_segment_pre - pre_extended - 1,
+                            total_samples_pre + num_samples_segment_pre + post_extended,
+                        )
+                    )
+            # exclude samples at the border of the recordings
+            if total_samples_pre > 0:
+                exclude_intervals.append((total_samples_pre - pre_extended - 1, total_samples_pre + post_extended))
+
+        total_valid_samples = (
+            total_num_samples - np.sum([end - start for start, end in exclude_intervals]) - pre_extended - post_extended
+        )
+        self.total_samples = int(total_valid_samples) if total_samples == -1 else total_samples
+        assert len(desired_shape) == 2, "desired_shape should be 2D"
+        assert (
+            desired_shape[0] * desired_shape[1] == recording.get_num_channels()
+        ), f"The product of desired_shape dimensions should be the number of channels: {recording.get_num_channels()}"
+        self.desired_shape = desired_shape
+
+        start_frame = start_frame if start_frame is not None else 0
+        end_frame = end_frame if end_frame is not None else total_num_samples
+        assert end_frame > start_frame, "end_frame must be greater than start_frame"
+
+        sequential_generator_params = dict()
+        sequential_generator_params["steps_per_epoch"] = steps_per_epoch
+        sequential_generator_params["pre_frame"] = pre_frame
+        sequential_generator_params["post_frame"] = post_frame
+        sequential_generator_params["batch_size"] = batch_size
+        sequential_generator_params["start_frame"] = start_frame
+        sequential_generator_params["end_frame"] = end_frame
+        sequential_generator_params["total_samples"] = self.total_samples
+        sequential_generator_params["pre_post_omission"] = pre_post_omission
+
+        json_path = tempfile.mktemp(suffix=".json")
+        with open(json_path, "w") as f:
+            json.dump(sequential_generator_params, f)
+        super().__init__(json_path)
+
+        self._update_end_frame(total_num_samples)
+
+        # self.list_samples will exclude the border intervals on the concat recording
+        self.recording_concat = concatenate_recordings(recordings)
+        self.exclude_intervals = exclude_intervals
+        self._calculate_list_samples(total_num_samples, exclude_intervals=exclude_intervals)
+
+        self._kwargs = dict(
+            recordings=recordings,
+            pre_frame=pre_frame,
+            post_frame=post_frame,
+            pre_post_omission=pre_post_omission,
+            desired_shape=desired_shape,
+            batch_size=batch_size,
+            steps_per_epoch=steps_per_epoch,
+            start_frame=start_frame,
+            end_frame=end_frame,
+        )
+
+    # this is overridden to exclude samples from borders
+    def _calculate_list_samples(self, total_frame_per_movie, exclude_intervals=[]):
+        # We first cut if start and end frames are too close to the edges.
+        self.start_sample = np.max([self.pre_frame + self.pre_post_omission, self.start_frame])
+        self.end_sample = np.min(
+            [
+                self.end_frame,
+                total_frame_per_movie - 1 - self.post_frame - self.pre_post_omission,
+            ]
+        )
+
+        if (self.end_sample - self.start_sample + 1) < self.batch_size:
+            raise Exception(
+                "Not enough frames to construct one "
+                + str(self.batch_size)
+                + " frame(s) batch between "
+                + str(self.start_sample)
+                + " and "
+                + str(self.end_sample)
+                + " frame number."
+            )
+
+        # +1 to make sure end_samples is included
+        list_samples_all = np.arange(self.start_sample, self.end_sample + 1)
+
+        if len(exclude_intervals) > 0:
+            for start, end in exclude_intervals:
+                list_samples_all = list_samples_all[(list_samples_all <= start) | (list_samples_all >= end)]
+            self.list_samples = list_samples_all
+        else:
+            self.list_samples = list_samples_all
+
+        if self.randomize:
+            np.random.shuffle(self.list_samples)
+
+        # We cut the number of samples if asked to
+        if self.total_samples > 0 and self.total_samples < len(self.list_samples):
+            self.list_samples = self.list_samples[0 : self.total_samples]
+
+    def __getitem__(self, index):
+        # This is to ensure we are going through
+        # the entire data when steps_per_epoch<self.__len__
+        shuffle_indexes = self.generate_batch_indexes(index)  # this not used during inferenc
+
+        input_full = np.zeros(
+            [len(shuffle_indexes), self.desired_shape[0], self.desired_shape[1], self.pre_frame + self.post_frame],
+            dtype="float32",
+        )
+        output_full = np.zeros([len(shuffle_indexes), self.desired_shape[0], self.desired_shape[1], 1], dtype="float32")
+
+        for batch_index, frame_index in enumerate(shuffle_indexes):
+            X, Y = self.__data_generation__(frame_index)
+
+            input_full[batch_index, :, :, :] = X
+            output_full[batch_index, :, :, :] = Y
+
+        return input_full, output_full
+
+    def __data_generation__(self, index_frame):
+        "Generates data containing batch_size samples"
+
+        # We reorganize to follow true geometry of probe for convolution
+        input_full = np.zeros(
+            [1, self.desired_shape[0], self.desired_shape[1], self.pre_frame + self.post_frame], dtype="float32"
+        )
+        output_full = np.zeros([1, self.desired_shape[0], self.desired_shape[1], 1], dtype="float32")
+
+        start_frame = index_frame - self.pre_frame - self.pre_post_omission
+        end_frame = index_frame + self.post_frame + self.pre_post_omission + 1
+        full_traces = self.recording_concat.get_traces(start_frame=start_frame, end_frame=end_frame).astype("float32")
+
+        if full_traces.shape[0] == 0:
+            print(f"Error! {index_frame}-{start_frame}-{end_frame}", flush=True)
+        output_frame_index = self.pre_frame + self.pre_post_omission
+        mask = np.ones(len(full_traces), dtype=bool)
+        mask = np.ones(len(full_traces), dtype=bool)
+        mask[output_frame_index - 1 : output_frame_index + 2] = False
+
+        data_img_input = full_traces[mask]
+        data_img_output = full_traces[output_frame_index][np.newaxis, :]
+
+        # make 3d based on desired shape
+        data_input_3d = data_img_input.reshape((-1, self.desired_shape[0], self.desired_shape[1]))
+        data_output_3d = data_img_output.reshape((-1, self.desired_shape[0], self.desired_shape[1]))
+
+        input_full[0] = data_input_3d.swapaxes(0, 1).swapaxes(1, 2)
+        output_full[0] = data_output_3d.swapaxes(0, 1).swapaxes(1, 2)
+
+        return input_full, output_full
+
+    def reshape_output(self, output):
+        return output.squeeze().reshape(-1, self.recording.get_num_channels())
+
+
+from deepinterpolation.generator_collection import SequentialGenerator
+
+
+class SpikeInterfaceRecordingSegmentGenerator(SequentialGenerator):
+    """This generator is used when dealing with a SpikeInterface recording.
+    The desired shape controls the reshaping of the input data before convolutions."""
+
+    def __init__(
+        self,
+        recording_segment: BaseRecordingSegment,
+        start_frame: int,
+        end_frame: int,
+        pre_frame: int = 30,
+        post_frame: int = 30,
+        pre_post_omission: int = 1,
+        desired_shape: tuple = (192, 2),
+        batch_size: int = 100,
+        steps_per_epoch: int = 10,
+    ):
+        self.recording_segment = recording_segment
+        self.num_channels = int(desired_shape[0] * desired_shape[1])
+        assert len(desired_shape) == 2, "desired_shape should be 2D"
+        self.desired_shape = desired_shape
+
+        num_segment_samples = recording_segment.get_num_samples()
+        start_frame = start_frame if start_frame is not None else 0
+        end_frame = end_frame if end_frame is not None else num_segment_samples
+        assert end_frame > start_frame, "end_frame must be greater than start_frame"
+        self.total_samples = end_frame - start_frame
+
+        sequential_generator_params = dict()
+        sequential_generator_params["steps_per_epoch"] = steps_per_epoch
+        sequential_generator_params["pre_frame"] = pre_frame
+        sequential_generator_params["post_frame"] = post_frame
+        sequential_generator_params["batch_size"] = batch_size
+        sequential_generator_params["start_frame"] = start_frame
+        sequential_generator_params["end_frame"] = end_frame
+        sequential_generator_params["total_samples"] = self.total_samples
+        sequential_generator_params["pre_post_omission"] = pre_post_omission
+
+        json_path = tempfile.mktemp(suffix=".json")
+        with open(json_path, "w") as f:
+            json.dump(sequential_generator_params, f)
+        super().__init__(json_path)
+
+        self._update_end_frame(num_segment_samples)
+        # IMPORTANT: this is used for inference, so we don't want to shuffle
+        self.randomize = False
+        self._calculate_list_samples(num_segment_samples)
+
+    def __getitem__(self, index):
+        # This is to ensure we are going through
+        # the entire data when steps_per_epoch<self.__len__
+        shuffle_indexes = self.generate_batch_indexes(index)  # this not used during inferenc
+
+        input_full = np.zeros(
+            [len(shuffle_indexes), self.desired_shape[0], self.desired_shape[1], self.pre_frame + self.post_frame],
+            dtype="float32",
+        )
+        output_full = np.zeros([len(shuffle_indexes), self.desired_shape[0], self.desired_shape[1], 1], dtype="float32")
+
+        for batch_index, frame_index in enumerate(shuffle_indexes):
+            X, Y = self.__data_generation__(frame_index)
+
+            input_full[batch_index, :, :, :] = X
+            output_full[batch_index, :, :, :] = Y
+
+        return input_full, output_full
+
+    def __data_generation__(self, index_frame):
+        "Generates data containing batch_size samples"
+
+        # We reorganize to follow true geometry of probe for convolution
+        input_full = np.zeros(
+            [1, self.desired_shape[0], self.desired_shape[1], self.pre_frame + self.post_frame], dtype="float32"
+        )
+        output_full = np.zeros([1, self.desired_shape[0], self.desired_shape[1], 1], dtype="float32")
+
+        start_frame = index_frame - self.pre_frame - self.pre_post_omission
+        end_frame = index_frame + self.post_frame + self.pre_post_omission + 1
+        channel_indices = slice(None)
+        full_traces = self.recording_segment.get_traces(
+            start_frame=start_frame, end_frame=end_frame, channel_indices=channel_indices
+        ).astype("float32")
+
+        if full_traces.shape[0] == 0:
+            print(f"Error! {index_frame}-{start_frame}-{end_frame}", flush=True)
+        output_frame_index = self.pre_frame + self.pre_post_omission
+        mask = np.ones(len(full_traces), dtype=bool)
+        mask = np.ones(len(full_traces), dtype=bool)
+        mask[output_frame_index - 1 : output_frame_index + 2] = False
+
+        data_img_input = full_traces[mask]
+        data_img_output = full_traces[output_frame_index][np.newaxis, :]
+
+        # make 3d based on desired shape
+        data_input_3d = data_img_input.reshape((-1, self.desired_shape[0], self.desired_shape[1]))
+        data_output_3d = data_img_output.reshape((-1, self.desired_shape[0], self.desired_shape[1]))
+
+        input_full[0] = data_input_3d.swapaxes(0, 1).swapaxes(1, 2)
+        output_full[0] = data_output_3d.swapaxes(0, 1).swapaxes(1, 2)
+
+        return input_full, output_full
+
+    def reshape_output(self, output):
+        return output.squeeze().reshape(-1, self.num_channels)
diff --git a/src/spikeinterface/preprocessing/deepinterpolation/tests/test_deepinterpolation.py b/src/spikeinterface/preprocessing/deepinterpolation/tests/test_deepinterpolation.py
new file mode 100644
index 0000000000..0525cdfc7a
--- /dev/null
+++ b/src/spikeinterface/preprocessing/deepinterpolation/tests/test_deepinterpolation.py
@@ -0,0 +1,178 @@
+import pytest
+import numpy as np
+from pathlib import Path
+
+import probeinterface as pi
+from spikeinterface import download_dataset, generate_recording, append_recordings, concatenate_recordings
+from spikeinterface.extractors import read_mearec, read_spikeglx, read_openephys
+from spikeinterface.preprocessing import depth_order, zscore
+
+from spikeinterface.preprocessing.deepinterpolation import train_deepinterpolation, deepinterpolate
+from spikeinterface.preprocessing.deepinterpolation import train_deepinterpolation, deepinterpolate
+
+
+try:
+    import tensorflow
+    import deepinterpolation
+
+    HAVE_DEEPINTERPOLATION = True
+except ImportError:
+    HAVE_DEEPINTERPOLATION = False
+
+
+if hasattr(pytest, "global_test_folder"):
+    cache_folder = pytest.global_test_folder / "deepinterpolation"
+else:
+    cache_folder = Path("cache_folder") / "deepinterpolation"
+
+
+def recording_and_shape():
+    num_cols = 2
+    num_rows = 64
+    probe = pi.generate_multi_columns_probe(num_columns=num_cols, num_contact_per_column=num_rows)
+    probe.set_device_channel_indices(np.arange(num_cols * num_rows))
+    recording = generate_recording(num_channels=num_cols * num_rows, durations=[10.0], sampling_frequency=30000)
+    recording.set_probe(probe, in_place=True)
+    recording = depth_order(recording)
+    recording = zscore(recording)
+    desired_shape = (num_rows, num_cols)
+    return recording, desired_shape
+
+
+@pytest.fixture
+def recording_and_shape_fixture():
+    return recording_and_shape()
+
+
+@pytest.mark.skipif(not HAVE_DEEPINTERPOLATION, reason="requires deepinterpolation")
+def test_deepinterpolation_generator_borders(recording_and_shape_fixture):
+    """Test that the generator avoids borders in multi-segment and recording lists cases"""
+    from spikeinterface.preprocessing.deepinterpolation.generators import SpikeInterfaceRecordingGenerator
+
+    recording, desired_shape = recording_and_shape_fixture
+    recording_multi_segment = append_recordings([recording, recording, recording])
+    recording_list = [recording, recording, recording]
+    recording_multi_list = [recording_multi_segment, recording_multi_segment, recording_multi_segment]
+
+    gen_multi_segment = SpikeInterfaceRecordingGenerator(recording_multi_segment, desired_shape=desired_shape)
+    gen_list = SpikeInterfaceRecordingGenerator(recording_list, desired_shape=desired_shape)
+    gen_multi_list = SpikeInterfaceRecordingGenerator(recording_multi_list, desired_shape=desired_shape)
+
+    # exclude in between segments
+    assert len(gen_multi_segment.exclude_intervals) == 2
+    # exclude in between recordings
+    assert len(gen_list.exclude_intervals) == 2
+    # exclude in between recordings and segments
+    assert len(gen_multi_list.exclude_intervals) == 2 * len(recording_multi_list) + 2
+
+
+@pytest.mark.skipif(not HAVE_DEEPINTERPOLATION, reason="requires deepinterpolation")
+def test_deepinterpolation_training(recording_and_shape_fixture):
+    recording, desired_shape = recording_and_shape_fixture
+
+    model_folder = Path(cache_folder) / "training"
+    # train
+    model_path = train_deepinterpolation(
+        recording,
+        model_folder=model_folder,
+        model_name="training",
+        train_start_s=1,
+        train_end_s=10,
+        train_duration_s=0.02,
+        test_start_s=0,
+        test_end_s=1,
+        test_duration_s=0.01,
+        pre_frame=20,
+        post_frame=20,
+        run_uid="si_test",
+        pre_post_omission=1,
+        desired_shape=desired_shape,
+    )
+    print(model_path)
+
+
+@pytest.mark.skipif(not HAVE_DEEPINTERPOLATION, reason="requires deepinterpolation")
+@pytest.mark.dependency(depends=["test_deepinterpolation_training"])
+def test_deepinterpolation_transfer(recording_and_shape_fixture, tmp_path):
+    recording, desired_shape = recording_and_shape_fixture
+
+    existing_model_path = Path(cache_folder) / "training" / "si_test_training_model.h5"
+    model_folder = Path(tmp_path) / "transfer"
+
+    # train
+    model_path = train_deepinterpolation(
+        recording,
+        model_folder=model_folder,
+        model_name="si_test_transfer",
+        existing_model_path=existing_model_path,
+        train_start_s=1,
+        train_end_s=10,
+        train_duration_s=0.02,
+        test_start_s=0,
+        test_end_s=1,
+        test_duration_s=0.01,
+        pre_frame=20,
+        post_frame=20,
+        pre_post_omission=1,
+        desired_shape=desired_shape,
+    )
+    print(model_path)
+
+
+@pytest.mark.skipif(not HAVE_DEEPINTERPOLATION, reason="requires deepinterpolation")
+@pytest.mark.dependency(depends=["test_deepinterpolation_training"])
+def test_deepinterpolation_inference(recording_and_shape_fixture):
+    recording, desired_shape = recording_and_shape_fixture
+    pre_frame = post_frame = 20
+    existing_model_path = Path(cache_folder) / "training" / "si_test_training_model.h5"
+
+    recording_di = deepinterpolate(
+        recording, model_path=existing_model_path, pre_frame=pre_frame, post_frame=post_frame, pre_post_omission=1
+    )
+    print(recording_di)
+    traces_original_first = recording.get_traces(start_frame=0, end_frame=100)
+    traces_di_first = recording_di.get_traces(start_frame=0, end_frame=100)
+    assert traces_di_first.shape == (100, recording.get_num_channels())
+    # first 20 frames should be the same
+    np.testing.assert_array_equal(traces_di_first[:pre_frame], traces_original_first[:pre_frame])
+    np.any(np.not_equal(traces_di_first[pre_frame:], traces_original_first[pre_frame:]))
+
+    num_samples = recording.get_num_samples()
+    traces_original_last = recording.get_traces(start_frame=num_samples - 100, end_frame=num_samples)
+    traces_di_last = recording_di.get_traces(start_frame=num_samples - 100, end_frame=num_samples)
+    # last 20 frames should be the same
+    np.testing.assert_array_equal(traces_di_last[-post_frame:], traces_original_last[-post_frame:])
+    np.any(np.not_equal(traces_di_last[:-post_frame:], traces_original_last[:-post_frame:]))
+
+
+@pytest.mark.skipif(not HAVE_DEEPINTERPOLATION, reason="requires deepinterpolation")
+@pytest.mark.dependency(depends=["test_deepinterpolation_training"])
+def test_deepinterpolation_inference_multi_job(recording_and_shape_fixture):
+    recording, desired_shape = recording_and_shape_fixture
+    pre_frame = post_frame = 20
+    existing_model_path = Path(cache_folder) / "training" / "si_test_training_model.h5"
+
+    recording_di = deepinterpolate(
+        recording,
+        model_path=existing_model_path,
+        pre_frame=pre_frame,
+        post_frame=post_frame,
+        pre_post_omission=1,
+        use_gpu=False,
+    )
+    print(recording_di)
+    recording_di_slice = recording_di.frame_slice(start_frame=0, end_frame=int(0.5 * recording.sampling_frequency))
+
+    recording_di_slice.save(folder=Path(cache_folder) / "di_slice", n_jobs=2, chunk_duration="50ms")
+    traces_chunks = recording_di_slice.get_traces()
+    traces_full = recording_di_slice.get_traces()
+    np.testing.assert_array_equal(traces_chunks, traces_full)
+
+
+if __name__ == "__main__":
+    recording_shape = recording_and_shape()
+    test_deepinterpolation_training(recording_shape)
+    # test_deepinterpolation_transfer()
+    test_deepinterpolation_inference(recording_shape)
+    # test_deepinterpolation_inference_multi_job()
+    # test_deepinterpolation_generator_borders(recording_shape)
diff --git a/src/spikeinterface/preprocessing/deepinterpolation/tf_utils.py b/src/spikeinterface/preprocessing/deepinterpolation/tf_utils.py
new file mode 100644
index 0000000000..277148e56b
--- /dev/null
+++ b/src/spikeinterface/preprocessing/deepinterpolation/tf_utils.py
@@ -0,0 +1,40 @@
+import os
+
+
+def has_tf(use_gpu=True, disable_tf_logger=True, memory_gpu=None):
+    try:
+        import_tf(use_gpu, disable_tf_logger, memory_gpu)
+        return True
+    except ImportError:
+        return False
+
+
+def import_tf(use_gpu=True, disable_tf_logger=True, memory_gpu=None):
+    if not use_gpu:
+        os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
+    if disable_tf_logger:
+        os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+    import tensorflow as tf
+
+    if disable_tf_logger:
+        tf.get_logger().setLevel("ERROR")
+
+    tf.compat.v1.disable_eager_execution()
+    gpus = tf.config.list_physical_devices("GPU")
+    if gpus and use_gpu:
+        if memory_gpu is None:
+            try:
+                # Currently, memory growth needs to be the same across GPUs
+                for gpu in gpus:
+                    tf.config.experimental.set_memory_growth(gpu, True)
+            except RuntimeError as e:
+                # Memory growth must be set before GPUs have been initialized
+                print(e)
+        else:
+            for gpu in gpus:
+                tf.config.set_logical_device_configuration(
+                    gpus[0], [tf.config.LogicalDeviceConfiguration(memory_limit=memory_gpu)]
+                )
+    return tf
diff --git a/src/spikeinterface/preprocessing/deepinterpolation/train.py b/src/spikeinterface/preprocessing/deepinterpolation/train.py
new file mode 100644
index 0000000000..c3b63dfc7e
--- /dev/null
+++ b/src/spikeinterface/preprocessing/deepinterpolation/train.py
@@ -0,0 +1,307 @@
+from __future__ import annotations
+import json
+import os
+import warnings
+from pathlib import Path
+from typing import Optional
+
+from concurrent.futures import ProcessPoolExecutor
+import multiprocessing as mp
+
+from .tf_utils import import_tf
+
+# from .generators import define_recording_generator_class
+
+from ...core import BaseRecording
+
+
+global train_func
+
+
+def train_deepinterpolation(
+    recordings: BaseRecording | list[BaseRecording],
+    model_folder: str | Path,
+    model_name: str,
+    desired_shape: tuple[int, int],
+    train_start_s: Optional[float] = None,
+    train_end_s: Optional[float] = None,
+    train_duration_s: Optional[float] = None,
+    test_start_s: Optional[float] = None,
+    test_end_s: Optional[float] = None,
+    test_duration_s: Optional[float] = None,
+    test_recordings: Optional[BaseRecording | list[BaseRecording]] = None,
+    pre_frame: int = 30,
+    post_frame: int = 30,
+    pre_post_omission: int = 1,
+    existing_model_path: Optional[str | Path] = None,
+    verbose: bool = True,
+    nb_gpus: int = 1,
+    steps_per_epoch: int = 10,
+    period_save: int = 100,
+    apply_learning_decay: int = 0,
+    nb_times_through_data: int = 1,
+    learning_rate: float = 0.0001,
+    loss: str = "mean_squared_error",
+    nb_workers: int = -1,
+    caching_validation: bool = False,
+    run_uid: str = "si",
+    network_name: str = "unet_single_ephys_1024",
+    use_gpu: bool = True,
+    disable_tf_logger: bool = True,
+    memory_gpu: Optional[int] = None,
+):
+    """
+    Train a deepinterpolation model from a recording extractor.
+
+    Parameters
+    ----------
+    recordings : BaseRecording | list[BaseRecording]
+        The recording(s) to be deepinteprolated. If a list is given, the recordings are concatenated
+        and samples at the border of the recordings are omitted.
+    model_folder : str | Path
+        Path to the folder where the model will be saved
+    model_name : str
+        Name of the model to be used
+    train_start_s : float or None, default: None
+        Start time of the training in seconds. If None, the training starts at the beginning of the recording
+    train_end_s : float or None, default: None
+        End time of the training in seconds. If None, the training ends at the end of the recording
+    train_duration_s : float, default: None
+        Duration of the training in seconds. If None, the entire [train_start_s, train_end_s] is used for training
+    test_start_s : float or None, default: None
+        Start time of the testing in seconds. If None, the testing starts at the beginning of the recording
+    test_end_s : float or None, default: None
+        End time of the testing in seconds. If None, the testing ends at the end of the recording
+    test_duration_s : float or None, default: None
+        Duration of the testing in seconds, If None, the entire [test_start_s, test_end_s] is used for testing (not recommended)
+    test_recordings : BaseRecording | list[BaseRecording] | None, default: None
+        The recording(s) used for testing. If None, the training recording (or recordings) is used for testing
+    desired_shape : tuple
+        Shape of the input to the network
+    pre_frame : int
+        Number of frames before the frame to be predicted
+    post_frame : int
+        Number of frames after the frame to be predicted
+    pre_post_omission : int
+        Number of frames to be omitted before and after the frame to be predicted
+    existing_model_path : str | Path
+        Path to an existing model to be used for transfer learning, default is None
+    verbose : bool
+        Whether to print the progress of the training, default is True
+    steps_per_epoch : int
+        Number of steps per epoch
+    period_save : int
+        Period of saving the model
+    apply_learning_decay : int
+        Whether to use a learning scheduler during training
+    nb_times_through_data : int
+        Number of times the data is repeated during training
+    learning_rate : float
+        Learning rate
+    loss : str
+        Loss function to be used
+    nb_workers : int
+        Number of workers to be used for the training
+    caching_validation : bool
+        Whether to cache the validation data, default is False
+    run_uid : str
+        Unique identifier for the training
+    network_name : str
+        Name of the network to be used, default is None
+    use_gpu : bool
+        Whether to use GPU, default is True
+    disable_tf_logger : bool
+        Whether to disable the tensorflow logger, default is True
+    memory_gpu : int
+        Amount of memory to be used by the GPU, default is None
+
+    Returns
+    -------
+    model_path : Path
+        Path to the model
+    """
+
+    if nb_workers == -1:
+        nb_workers = os.cpu_count()
+
+    args = (
+        recordings,
+        model_folder,
+        model_name,
+        desired_shape,
+        train_start_s,
+        train_end_s,
+        train_duration_s,
+        test_start_s,
+        test_end_s,
+        test_duration_s,
+        test_recordings,
+        pre_frame,
+        post_frame,
+        pre_post_omission,
+        existing_model_path,
+        verbose,
+        nb_gpus,
+        steps_per_epoch,
+        period_save,
+        apply_learning_decay,
+        nb_times_through_data,
+        learning_rate,
+        loss,
+        nb_workers,
+        caching_validation,
+        run_uid,
+        network_name,
+        use_gpu,
+        disable_tf_logger,
+        memory_gpu,
+    )
+    global train_func
+    train_func = train_deepinterpolation_process
+    with ProcessPoolExecutor(mp_context=mp.get_context("spawn"), max_workers=nb_workers) as executor:
+        f = executor.submit(train_func, *args)
+        model_path = f.result()
+    return model_path
+
+
+def train_deepinterpolation_process(
+    recordings: BaseRecording | list[BaseRecording],
+    model_folder: str | Path,
+    model_name: str,
+    desired_shape: tuple[int, int],
+    train_start_s: float,
+    train_end_s: float,
+    train_duration_s: float | None,
+    test_start_s: float,
+    test_end_s: float,
+    test_duration_s: float | None,
+    test_recordings: Optional[BaseRecording | list[BaseRecording]] = None,
+    pre_frame: int = 30,
+    post_frame: int = 30,
+    pre_post_omission: int = 1,
+    existing_model_path: Optional[str | Path] = None,
+    verbose: bool = True,
+    nb_gpus: int = 1,
+    steps_per_epoch: int = 10,
+    period_save: int = 100,
+    apply_learning_decay: int = 0,
+    nb_times_through_data: int = 1,
+    learning_rate: float = 0.0001,
+    loss: str = "mean_absolute_error",
+    nb_workers: int = -1,
+    caching_validation: bool = False,
+    run_uid: str = "training",
+    network_name: str = "unet_single_ephys_1024",
+    use_gpu: bool = True,
+    disable_tf_logger: bool = True,
+    memory_gpu: Optional[int] = None,
+):
+    from deepinterpolation.trainor_collection import core_trainer
+    from deepinterpolation.generic import ClassLoader
+    from .generators import SpikeInterfaceRecordingGenerator
+
+    # initialize TF
+    _ = import_tf(use_gpu, disable_tf_logger, memory_gpu=memory_gpu)
+
+    trained_model_folder = Path(model_folder)
+    trained_model_folder.mkdir(exist_ok=True, parents=True)
+
+    # check if roughly zscored
+    if not isinstance(recordings, list):
+        recordings = [recordings]
+    fs = recordings[0].sampling_frequency
+
+    ### Define params
+    start_frame_training = int(train_start_s * fs)
+    end_frame_training = int(train_end_s * fs)
+    if train_duration_s is not None:
+        total_samples_training = int(train_duration_s * fs)
+    else:
+        total_samples_training = -1
+    start_frame_test = int(test_start_s * fs)
+    end_frame_test = int(test_end_s * fs)
+    if test_duration_s is not None:
+        total_samples_testing = int(test_duration_s * fs)
+    else:
+        total_samples_testing = -1
+
+    if test_recordings is None:
+        test_recordings = recordings
+        if (start_frame_training <= start_frame_test < end_frame_training) or (
+            start_frame_training < end_frame_test <= end_frame_training
+        ):
+            warnings.warn("Training and testing overlap. This is not recommended.")
+
+    # Those are parameters used for the network topology
+    network_params = dict()
+    network_params["type"] = "network"
+    # Name of network topology in the collection
+    network_params["name"] = network_name if network_name is not None else "unet_single_ephys_1024"
+    training_params = dict()
+    training_params["output_dir"] = str(trained_model_folder)
+    # We pass on the uid
+    training_params["run_uid"] = run_uid
+
+    # We convert to old schema
+    training_params["nb_gpus"] = nb_gpus
+    training_params["type"] = "trainer"
+    training_params["steps_per_epoch"] = steps_per_epoch
+    training_params["period_save"] = period_save
+    training_params["apply_learning_decay"] = apply_learning_decay
+    training_params["nb_times_through_data"] = nb_times_through_data
+    training_params["learning_rate"] = learning_rate
+    training_params["loss"] = loss
+    training_params["nb_workers"] = nb_workers
+    training_params["caching_validation"] = caching_validation
+    training_params["model_string"] = model_name
+    if existing_model_path:
+        training_params["model_path"] = str(existing_model_path)
+
+    # Training (from core_trainor class)
+    training_data_generator = SpikeInterfaceRecordingGenerator(
+        recordings,
+        pre_frame=pre_frame,
+        post_frame=post_frame,
+        pre_post_omission=pre_post_omission,
+        start_frame=start_frame_training,
+        end_frame=end_frame_training,
+        desired_shape=desired_shape,
+        total_samples=total_samples_training,
+    )
+    test_data_generator = SpikeInterfaceRecordingGenerator(
+        test_recordings,
+        pre_frame=pre_frame,
+        post_frame=post_frame,
+        pre_post_omission=pre_post_omission,
+        start_frame=start_frame_test,
+        end_frame=end_frame_test,
+        steps_per_epoch=-1,
+        desired_shape=desired_shape,
+        total_samples=total_samples_testing,
+    )
+
+    network_json_path = trained_model_folder / "network_params.json"
+    with open(network_json_path, "w") as f:
+        json.dump(network_params, f)
+
+    network_obj = ClassLoader(network_json_path)
+    data_network = network_obj.find_and_build()(network_json_path)
+
+    training_json_path = trained_model_folder / "training_params.json"
+    with open(training_json_path, "w") as f:
+        json.dump(training_params, f)
+
+    training_class = core_trainer(training_data_generator, test_data_generator, data_network, training_json_path)
+
+    if verbose:
+        print("Created objects for training. Running training job")
+    training_class.run()
+
+    if verbose:
+        print("Training job finished - finalizing output model")
+    training_class.finalize()
+
+    # Re-load model from output folder
+    model_path = trained_model_folder / f"{training_params['run_uid']}_{training_params['model_string']}_model.h5"
+
+    return model_path
diff --git a/src/spikeinterface/preprocessing/preprocessinglist.py b/src/spikeinterface/preprocessing/preprocessinglist.py
index de83d509a8..67dcd7be45 100644
--- a/src/spikeinterface/preprocessing/preprocessinglist.py
+++ b/src/spikeinterface/preprocessing/preprocessinglist.py
@@ -29,7 +29,7 @@
 from .silence_periods import SilencedPeriodsRecording, silence_periods
 from .phase_shift import PhaseShiftRecording, phase_shift
 from .zero_channel_pad import ZeroChannelPaddedRecording, zero_channel_pad
-from .deepinterpolation import DeepInterpolatedRecording, deepinterpolate
+from .deepinterpolation import DeepInterpolatedRecording, deepinterpolate, train_deepinterpolation
 from .highpass_spatial_filter import HighpassSpatialFilterRecording, highpass_spatial_filter
 from .interpolate_bad_channels import InterpolateBadChannelsRecording, interpolate_bad_channels
 from .average_across_direction import AverageAcrossDirectionRecording, average_across_direction