diff --git a/grain/_src/python/dataset/transformations/process_prefetch_test.py b/grain/_src/python/dataset/transformations/process_prefetch_test.py
index 610cb7eb2..fd97d3e80 100644
--- a/grain/_src/python/dataset/transformations/process_prefetch_test.py
+++ b/grain/_src/python/dataset/transformations/process_prefetch_test.py
@@ -14,6 +14,7 @@
 from collections.abc import Sequence
 from concurrent import futures
 import dataclasses
+import gc
 import logging as std_logging
 import os
 import sys