From f26a7cc082d427e9cce22b8c7558facf9f854f07 Mon Sep 17 00:00:00 2001
From: Lena Kashtelyan <drfreund@fb.com>
Date: Tue, 14 Sep 2021 18:20:20 -0700
Subject: [PATCH] Add `batch_number` back into SAAABO NEHVI tutorial

Summary:
As titled

Created from CodeHub with https://fburl.com/edit-in-codehub

Reviewed By: dme65

Differential Revision: D30943852

fbshipit-source-id: 2170bcadab19c7ef99a4f3d496c88da77a630777
---
 tutorials/saasbo_nehvi.ipynb | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tutorials/saasbo_nehvi.ipynb b/tutorials/saasbo_nehvi.ipynb
index cdb948a581a..dff53610676 100644
--- a/tutorials/saasbo_nehvi.ipynb
+++ b/tutorials/saasbo_nehvi.ipynb
@@ -1038,7 +1038,8 @@
         "\n",
         "n_results = N_BATCH*BATCH_SIZE + N_INIT\n",
         "\n",
-        "sc = axes.scatter(train_obj[:, 0], train_obj[:,1], c=df.trial_index.values, alpha=0.8)\n",
+        "batch_number = df.trial_index.values\n",
+        "sc = axes.scatter(train_obj[:, 0], train_obj[:,1], c=batch_number, alpha=0.8)\n",
         "axes.set_title(algos[0])\n",
         "axes.set_xlabel(\"Objective 1\")\n",
         "axes.set_ylabel(\"Objective 2\")\n",