put these changes to another PR.

python/dgl/graphbolt/impl/feature_cache.py

@@ -1,7 +1,7 @@
 """CPU Feature Cache implementation wrapper for graphbolt."""
 import torch
 
-__all__ = ["CPUFeatureCache"]
+__all__ = ["FeatureCache"]
 
 caching_policies = {
     "s3-fifo": torch.ops.graphbolt.s3_fifo_cache_policy,
@@ -11,7 +11,7 @@
 }
 
 
-class CPUFeatureCache(object):
+class FeatureCache(object):
     r"""High level wrapper for the CPU feature cache.
 
     Parameters
@@ -34,12 +34,10 @@ def __init__(
         self,
         cache_shape,
         dtype,
-        policy=None,
+        policy="sieve",
         num_parts=None,
         pin_memory=False,
     ):
-        if policy is None:
-            policy = "sieve"
         assert (
             policy in caching_policies
         ), f"{list(caching_policies.keys())} are the available caching policies."

tests/python/pytorch/graphbolt/impl/test_feature_cache.py

@@ -29,7 +29,7 @@ def test_feature_cache(dtype, feature_size, num_parts, policy):
         torch.get_num_threads() if num_parts is None else num_parts
     )
     a = torch.randint(0, 2, [1024, feature_size], dtype=dtype)
-    cache = gb.impl.CPUFeatureCache(
+    cache = gb.impl.FeatureCache(
         (cache_size,) + a.shape[1:], a.dtype, policy, num_parts
     )