From 0c336d64575f01107f49c5522b9d454d00fcc7d3 Mon Sep 17 00:00:00 2001
From: Ebrahim Songhori <esonghori@google.com>
Date: Thu, 16 Nov 2023 11:45:18 -0800
Subject: [PATCH] Left from previous fix.

PiperOrigin-RevId: 583115192
Change-Id: I2b9609ac5ef7881cc442bbfd239ace63f6894453
---
 tools/docker/ubuntu_tf_agents | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tools/docker/ubuntu_tf_agents b/tools/docker/ubuntu_tf_agents
index 8071ce38f..fb443a593 100644
--- a/tools/docker/ubuntu_tf_agents
+++ b/tools/docker/ubuntu_tf_agents
@@ -87,7 +87,8 @@ RUN $python_version -m pip install wheel \
       protobuf==4.21.12 \
       google-api-core \
       google-cloud-bigquery \
-      psutil
+      psutil \
+      tf-keras
 
 ######################################################
 # Installs Jupyter notebook.