diff --git a/src/djctools/module_extensions.py b/src/djctools/module_extensions.py
index 6652d87..617175e 100644
--- a/src/djctools/module_extensions.py
+++ b/src/djctools/module_extensions.py
@@ -236,6 +236,8 @@ class PlottingModule(torch.nn.Module):
     """
     This layer is used to enable or disable plotting from within the model.
     It is meant as a base class from which to inherit, and should not be used directly.
+    The inhereting classes must not return anythin from the plot function, otherwise issues with model saving and
+    multi-GPU training can occur!!
     The logic works as follows:
       - If plotting is enabled, the forward method caches the data given to it while the model is executing. 
         It does not return anything, and it also does not start any plotting process.
@@ -323,13 +325,14 @@ def _plot_worker(self):
         self._cache = []  # Clear the cache before plotting
         self.plot(data)
 
-    def plot(self, data):
+    def plot(self, data)-> None:
         """
         Override this method in the subclass to implement custom plotting logic.
+        !!The ploting functions must not return anything, otherwise model saving and mult-GPU training will not work correctly!!
         Args:
             data (list): Cached data to be plotted.
         """
-        raise NotImplementedError("The 'plot' method must be implemented in subclasses.")
+        raise NotImplementedError("The 'plot' method must be implemented in subclasses. It is not allowed to return anything.")
 
     def _join_plot_thread(self):
         """
diff --git a/src/djctools/training.py b/src/djctools/training.py
index 16f4f91..41a76b6 100644
--- a/src/djctools/training.py
+++ b/src/djctools/training.py
@@ -1,7 +1,7 @@
 # import as from djctools.training
 
 import torch
-from .module_extensions import flush_all_plotting
+from .module_extensions import flush_all_plotting, PlottingModule
 from .wandb_tools import wandb_wrapper
 import numpy as np
 import os
@@ -93,6 +93,7 @@ def __init__(self, model, optimizer, num_gpus=1, device_ids=None, verbose_level=
             self.device_ids = device_ids if device_ids is not None else list(range(num_gpus))
             self.device = f'cuda:{self.device_ids[0]}'
             self.devices = [f'cuda:{device_id}' for device_id in self.device_ids]
+
         else:
             # Fall back to CPU if no GPU available or num_gpus is 1
             self.device = 'cpu'
@@ -115,10 +116,14 @@ def __init__(self, model, optimizer, num_gpus=1, device_ids=None, verbose_level=
         
         print(f"Creating replicas using devices: {self.devices}")
 
+        #TODO: fix plotting modules for multi-GPU training, currently just set to None
+        if num_gpus > 1:
+            model = self.remove_plotting(model) #set all plotting modules to None to avoid deepcopy issues during multi-GPU training
         self.model_replicas = make_replicas(model, self.devices)
         self.optimizer = optimizer
         self.verbose_level = verbose_level
 
+
     def save_model(self, filepath):
         """
         Saves the model (not just weights) to a file. 
@@ -128,6 +133,7 @@ def save_model(self, filepath):
         """
         torch.save(self.model_replicas[0], filepath) #the first replica is always the master
 
+
     def load_model(self, filepath):
         """
         Loads model from a file. slim wrapper
@@ -281,4 +287,20 @@ def val_batch_callback(self, model, batch_number, batch_data):
         and should be implemented by the user through inheritance. Please do not use for logging purposes,
         use the wandb_wrapper.log() function instead.
         """
-        pass
\ No newline at end of file
+        pass
+
+    def remove_plotting(self, model):
+        """
+        Sets all plotting modules in model to None to avoid deepcopy issues during multi-GPU training.
+        Plotting modules are not allowed to return anything from their plot function, otherwise this causes issues.
+
+        Args:
+        -----
+            model: the model from which to remove the plotting modules
+        """
+        print("Removing plotting modules for multi-GPU training...")
+        for name, module in model._modules.items():
+            if isinstance(module, PlottingModule):
+                model._modules[name] = None
+        
+        return model