Avoid internal comparisons of ProxyObjects

rapidsai · Dec 2, 2020 · f2fffb1 · f2fffb1
1 parent d0f596f
commit f2fffb1
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/dask_cuda/object_spilling_host_file.py b/dask_cuda/object_spilling_host_file.py
@@ -61,6 +61,7 @@ def unspilled_proxies(self):
         return ret
 
     def obj_mappings(self):
+        # TODO: simplify and optimize
         proxied_id_to_proxy = {}
         buffer_to_proxies = {}
 
@@ -71,7 +72,7 @@ def obj_mappings(self):
             proxied_id_to_proxy[proxied_id] = p
             for buf in get_device_memory_objects(proxied):
                 l = buffer_to_proxies.get(buf, [])
-                if p not in l:
+                if id(p) not in set(id(i) for i in l):
                     l.append(p)
                 buffer_to_proxies[buf] = l
         return proxied_id_to_proxy, buffer_to_proxies