Systems, apparatuses, and methods include technology that determines, with a neural network, that a first eviction node stored in a cache will be evicted from the cache based on a cache policy. The first eviction node is part of a plurality of nodes associated with a graph. Further, a subset of nodes of the plurality of nodes remains in the cache after the eviction of the first eviction node from the cache. The technology further tracks a number of cache hits on the cache during an aggregation operation associated with a hardware accelerator, where the aggregation operation is executed on the subset of nodes that remain in the cache after the eviction of the eviction node from the cache. The technology executes a training process on the neural network to adjust the cache policy based on the number of the cache hits.