Update input_preprocess.py

input_preprocess.py

@@ -23,7 +23,7 @@ def scenarios_list():
     ])
 
 #%% Token Generation
-def tokenizer(deepmimo_data, gen_raw=True):
+def tokenizer(deepmimo_data, scenario_idxs, gen_raw=True):
     """
     Generates tokens by preparing and preprocessing the dataset.
 
@@ -40,6 +40,7 @@ def tokenizer(deepmimo_data, gen_raw=True):
     """
     
     # Patch generation or loading
+    deepmimo_data = deepmimo_data[scenario_idxs]
     n_scenarios = len(deepmimo_data)
     patches = [patch_maker(deepmimo_data[scenario_idx]) for scenario_idx in range(n_scenarios)]
     patches = np.vstack(patches)