From f3f57330fb63ef80efb9ced50aaab5c366de4323 Mon Sep 17 00:00:00 2001 From: Muhammed Fatih BALIN Date: Sat, 20 Jul 2024 13:23:38 -0400 Subject: [PATCH] Update python/dgl/graphbolt/impl/gpu_cached_feature.py --- python/dgl/graphbolt/impl/gpu_cached_feature.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/dgl/graphbolt/impl/gpu_cached_feature.py b/python/dgl/graphbolt/impl/gpu_cached_feature.py index 362224bea0c3..2b81213d1b11 100644 --- a/python/dgl/graphbolt/impl/gpu_cached_feature.py +++ b/python/dgl/graphbolt/impl/gpu_cached_feature.py @@ -128,7 +128,9 @@ class _Waiter: @staticmethod def wait(): """Returns the stored value when invoked.""" - values[missing_index] = missing_values_future.wait() + missing_values = missing_values_future.wait() + self._feature.replace(missing_keys, missing_values) + values[missing_index] = missing_values return values yield _Waiter()