diff --git a/main.py b/main.py
index 86a8bcb..e7f8c98 100644
--- a/main.py
+++ b/main.py
@@ -103,14 +103,10 @@ def main_hyperband():
         "kernels_main": [1, 3, 5, 7, 9],
         "dense_main": [16, 32, 64, 128, 256, 512],
     }
-    param = hyperband.sample_params(params)
-    logger.info(param)
 
     logger.info("create training dataset")
-    domain_tr, flow_tr, client_tr, server_tr = dataset.create_dataset_from_flows(user_flow_df, char_dict,
-                                                                                 max_len=args.domain_length,
-                                                                                 window_size=args.window)
-
+    domain_tr, flow_tr, client_tr, server_tr = load_or_generate_h5data(args.h5data, args.train_data,
+                                                                       args.domain_length, args.window)
     hp = hyperband.Hyperband(params,
                              [domain_tr, flow_tr],
                              [client_tr, server_tr])