Reduced server GPU memory for Parallel Swarm Learning

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(),