Update inference.py

inference.py

@@ -33,7 +33,7 @@ if torch.cuda.is_available():
 # Folders
 # MODELS_FOLDER = 'models/'
 
-def dataset_gen(preprocessed_chs, input_type, scenario_idxs, lwm_model):
+def dataset_gen(preprocessed_chs, input_type, lwm_model):
     
     if input_type in ['cls_emb', 'channel_emb']:
         dataset = prepare_for_LWM(preprocessed_chs, device)
@@ -41,11 +41,6 @@ def dataset_gen(preprocessed_chs, input_type, scenario_idxs, lwm_model):
         dataset = create_raw_dataset(preprocessed_chs, device)
     
     if input_type in ['cls_emb','channel_emb']:
-        # model = LWM().to(device)
-        # ckpt_name = 'model_weights.pth'
-        # ckpt_path = os.path.join(MODELS_FOLDER, ckpt_name)
-        # lwm_model = load_model(model, ckpt_path, device)
-        # print(f"Model loaded successfully on {device}")
 
         # Process data through LWM
         lwm_loss, embedding_data = evaluate(lwm_model, dataset)