Upload 3 files

inference.py

@@ -30,10 +30,8 @@ device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
 if torch.cuda.is_available():
     torch.cuda.empty_cache()
 
-# Folders
-# MODELS_FOLDER = 'models/'
 
-def dataset_gen(preprocessed_chs, input_type, lwm_model):
+def lwm_inference(preprocessed_chs, input_type, lwm_model):
     
     if input_type in ['cls_emb', 'channel_emb']:
         dataset = prepare_for_LWM(preprocessed_chs, device)

lwm_model.py

@@ -11,9 +11,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
 from inference import *
-from load_data import load_DeepMIMO_data
 from input_preprocess import *
-from huggingface_hub import hf_hub_download
 
 
 ELEMENT_LENGTH = 16
@@ -131,8 +129,9 @@ class LWM(torch.nn.Module):
         model = cls().to(device)
 
         # Download model weights using Hugging Face Hub
-        ckpt_path = hf_hub_download(repo_id="sadjadalikhani/LWM", filename=ckpt_name, use_auth_token=use_auth_token)
-
+        # ckpt_path = hf_hub_download(repo_id="sadjadalikhani/LWM", filename=ckpt_name, use_auth_token=use_auth_token)
+        ckpt_path = ckpt_name
+        
         # Load the model weights
         model.load_state_dict(torch.load(ckpt_path, map_location=device))
         print(f"Model loaded successfully from {ckpt_path} to {device}")