Upload inference.py

inference.py

@@ -33,24 +33,18 @@ if torch.cuda.is_available():
 
 def lwm_inference(preprocessed_chs, input_type, lwm_model):
     
-    if input_type in ['cls_emb', 'channel_emb']:
-        dataset = prepare_for_LWM(preprocessed_chs, device)
-    elif input_type == 'raw':
-        dataset = create_raw_dataset(preprocessed_chs, device)
+    dataset = prepare_for_LWM(preprocessed_chs, device)
     
-    if input_type in ['cls_emb','channel_emb']:
-
-        # Process data through LWM
-        lwm_loss, embedding_data = evaluate(lwm_model, dataset)
-        
-        print(f'LWM loss: {lwm_loss:.4f}')
-        
-        if input_type == 'cls_emb':
-            embedding_data = embedding_data[:, 0]
-        elif input_type == 'channel_emb':  
-            embedding_data = embedding_data[:, 1:]
-        
-        dataset = embedding_data.float()
+    # Process data through LWM
+    lwm_loss, embedding_data = evaluate(lwm_model, dataset)
+    print(f'LWM loss: {lwm_loss:.4f}')
+    
+    if input_type == 'cls_emb':
+        embedding_data = embedding_data[:, 0]
+    elif input_type == 'channel_emb':  
+        embedding_data = embedding_data[:, 1:]
+    
+    dataset = embedding_data.float()
         
     return dataset