From e4024e01d569f9379c8a6f4db8d1032fc7d0fb76 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Sat, 21 Dec 2019 18:33:15 +0200
Subject: [PATCH] Make sure ProxyDictPreWrite and ProxyDictPostWrite are
 pickled correctly

---
 trains/utilities/proxy_object.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/trains/utilities/proxy_object.py b/trains/utilities/proxy_object.py
index bb63d6d2..d1cadd01 100644
--- a/trains/utilities/proxy_object.py
+++ b/trains/utilities/proxy_object.py
@@ -17,6 +17,9 @@ class ProxyDictPostWrite(dict):
         super(ProxyDictPostWrite, self).__setitem__(key, value)
         self._set_callback()
 
+    def __reduce__(self):
+        return dict, (), None, None, iter(self._to_dict().items())
+
     def _set_callback(self, *_):
         if self._update_func:
             self._update_func(self._update_obj, self)
@@ -48,6 +51,9 @@ class ProxyDictPreWrite(dict):
         self._update_obj = update_obj
         self._update_func = update_func
 
+    def __reduce__(self):
+        return dict, (), None, None, iter(self.items())
+
     def __setitem__(self, key, value):
         key_value = self._set_callback((key, value,))
         if key_value: