def gpu_cached_feature(
    fallback_features: Union[Feature, Dict[FeatureKey, Feature]],
    max_cache_size_in_bytes: int,
) -> Union[GPUCachedFeature, Dict[FeatureKey, GPUCachedFeature]]:
    r"""GPU cached feature wrapping a fallback feature. It uses the least
    recently used (LRU) algorithm as the cache eviction policy.

    Places the GPU cache to torch.cuda.current_device().

    Parameters
    ----------
    fallback_features : Union[Feature, Dict[FeatureKey, Feature]]
        The fallback feature(s).
    max_cache_size_in_bytes : int
        The capacity of the GPU cache in bytes.
    Returns
    -------
    Union[GPUCachedFeature, Dict[FeatureKey, GPUCachedFeature]]
        The feature(s) wrapped with GPUCachedFeature.
    """
    return wrap_with_cached_feature(
        GPUCachedFeature, fallback_features, max_cache_size_in_bytes
    )
GPU cached feature wrapping a fallback feature. It uses the least
    recently used (LRU) algorithm as the cache eviction policy.

    Places the GPU cache to torch.cuda.current_device().

    Parameters
    ----------
    fallback_features : Union[Feature, Dict[FeatureKey, Feature]]
        The fallback feature(s).
    max_cache_size_in_bytes : int
        The capacity of the GPU cache in bytes.
    Returns
    -------
    Union[GPUCachedFeature, Dict[FeatureKey, GPUCachedFeature]]
        The feature(s) wrapped with GPUCachedFeature.
    
python/dgl/graphbolt/impl/gpu_cached_feature.py