diff --git a/edml/core/server.py b/edml/core/server.py index a671ba98e404c6981b35eb90d3e5bb1c17084304..3c9833e5a965a7afd1d3061719c0c50ffdf536ce 100644 --- a/edml/core/server.py +++ b/edml/core/server.py @@ -335,7 +335,14 @@ class DeviceServer: self._lr_scheduler.step(round_no + 1) # epoch=1 else: self._lr_scheduler.step() - + # delete references and free GPU memory manually + server_batch = None + server_labels = None + server_gradients = None + client_gradients = None + concatenated_client_gradients = None + mean_tensor = None + torch.cuda.empty_cache() return ( self.node_device.client.get_weights(), self.get_weights(),