Update input_preprocess.py

input_preprocess.py

@@ -40,7 +40,7 @@ def tokenizer(deepmimo_data, scenario_idxs, gen_raw=True):
     """
     
     # Patch generation or loading
-    deepmimo_data = deepmimo_data[scenario_idxs]
+    deepmimo_data = [deepmimo_data[i] for i in 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)