From 19742636dd7f415b746968e3105a6a21dc243b07 Mon Sep 17 00:00:00 2001
From: Ian Stenbit <3072903+ianstenbit@users.noreply.github.com>
Date: Tue, 21 Mar 2023 16:31:46 -0600
Subject: [PATCH] Avoid caching the entire Waymo Open Dataset during loading
 (#1559)

* Avoid caching the entire Waymo Open Dataset during loading

* Formatting
---
 keras_cv/datasets/waymo/load.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/keras_cv/datasets/waymo/load.py b/keras_cv/datasets/waymo/load.py
index b186984a6d..cea58285d6 100644
--- a/keras_cv/datasets/waymo/load.py
+++ b/keras_cv/datasets/waymo/load.py
@@ -15,7 +15,6 @@
 import os
 
 import tensorflow as tf
-import tensorflow_datasets as tfds
 
 from keras_cv.datasets.waymo import transformer
 from keras_cv.utils import assert_waymo_open_dataset_installed
@@ -28,9 +27,9 @@
 
 def _generate_frames(segments, transformer):
     def _generator():
-        for record in tfds.as_numpy(segments):
+        for record in segments:
             frame = waymo_open_dataset.dataset_pb2.Frame()
-            frame.ParseFromString(record)
+            frame.ParseFromString(record.numpy())
             yield transformer(frame)
 
     return _generator