Set CPU as default location for unpickling received data to decrease GPU memory consumption

edml/helpers/proto_helpers.py

@@ -27,7 +27,8 @@ def _tensor_to_bytes(tensor: torch.Tensor) -> bytes:
 
 
 def _bytes_to_tensor(raw_bytes: bytes) -> torch.Tensor:
-    return pickle.loads(raw_bytes)
+    bs = io.BytesIO(raw_bytes)
+    return CpuUnpickler(bs).load()
 
 
 def _state_dict_to_bytes(state_dict: StateDict) -> bytes:
@@ -46,7 +47,8 @@ def _metrics_to_bytes(
 
 
 def _bytes_to_metrics(raw_bytes: bytes):
-    return pickle.loads(raw_bytes)
+    bs = io.BytesIO(raw_bytes)
+    return CpuUnpickler(bs).load()
 
 
 def tensor_to_proto(tensor: torch.Tensor) -> datastructures_pb2.Tensor: