Update python/dgl/graphbolt/impl/gpu_cached_feature.py

dmlc · Jul 20, 2024 · f3f5733 · f3f5733
1 parent 5625d3a
commit f3f5733
Showing 1 changed file with 3 additions and 1 deletion.
diff --git 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()