make the code more robust.

dmlc · Aug 21, 2024 · 0560739 · 0560739
1 parent c431578
commit 0560739
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/dgl/graphbolt/feature_store.py b/python/dgl/graphbolt/feature_store.py
@@ -338,8 +338,8 @@ def wrap_with_cached_feature(
         cache_size = bytes_to_number_of_items(max_cache_size_in_bytes, feat0)
         if cache is None:
             cache = cached_feature_type._cache_type(
-                (cache_size,) + feat0.shape[1:],
-                feat0.dtype,
+                cache_shape=(cache_size,) + feat0.shape[1:],
+                dtype=feat0.dtype,
                 *args,
                 **kwargs,
             )