From a1162cf83279e2fe566fda4a955a291ec7d96557 Mon Sep 17 00:00:00 2001
From: stannis <stannis@vicarious.com>
Date: Wed, 1 Dec 2021 11:22:42 -0800
Subject: [PATCH] Revert default train size to 20

---
 examples/rcn.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/examples/rcn.py b/examples/rcn.py
index 7768f3b1..3a37bbb7 100644
--- a/examples/rcn.py
+++ b/examples/rcn.py
@@ -39,7 +39,7 @@
 
 # Use train_size = 100 if you have a gpu with atleast 8Gbs memory.
 # Recommend that jax is installed with cuda enabled for this option.
-train_size = 100
+train_size = 20
 test_size = 20
 
 
@@ -98,7 +98,9 @@ def fetch_mnist_dataset(test_size: int, seed: int = 5):
 # # 2. Load the model
 
 # %%
-data = np.load("example_data/rcn_100.npz", allow_pickle=True, encoding="latin1")
+data = np.load(
+    f"example_data/rcn_{train_size}.npz", allow_pickle=True, encoding="latin1"
+)
 frcs, edges, suppression_masks, filters = (
     data["frcs"][:train_size],
     data["edges"][:train_size],