“Se ha producido un error de tiempo de espera de Google Drive” al cargar una carpeta con una gran cantidad de archivos pequeños en Google Colab.

Estoy trabajando en una carpeta con 150 mil archivos (~64 kb de tamaño cada uno) almacenados en mi Google Drive. En Google Colab, estoy usando la siguiente función para obtener la lista de todos los archivos. Colab arroja el error de tiempo de espera de la unidad cuando ejecuté la función en la carpeta con 150k archivos. Intenté volver a montar la unidad varias veces, reiniciar Colab y usar glob para obtener la lista de archivos. Ninguno de estos enfoques funcionó para mí.

¿Sabes cómo puedo trabajar efectivamente con 150k archivos?

def list_all_files_with_an_extension(pth_folder: str, file_extension: list = [".png"]) -> list:
    """
    Obtener archivos de una carpeta con una extensión
    """
    # Espacio en espera
    list_files = []
    # Recorrer la carpeta y subcarpetas
    for root, dirs, files in os.walk(pth_folder):
        for file in files:
            for extension in file_extension:
                if file.endswith(extension):
                    list_files.append(os.path.join(root, file))
    # Devuelve la lista de rutas
    return list_files