Update inference.py

inference.py

@@ -23,19 +23,18 @@ import numpy as np
 import warnings
 warnings.filterwarnings('ignore')
 
-# Set random seeds for reproducibility across CPU and GPU
-def set_random_seed(seed=42):
+def set_seed(seed=42):
     torch.manual_seed(seed)
     np.random.seed(seed)
-    #random.seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(seed)
-    # Ensures deterministic behavior
-    torch.backends.cudnn.deterministic = True
-    torch.backends.cudnn.benchmark = False
 
-# Apply random seed
-set_random_seed()
+# Use this function at the start of your code
+set_seed(42)
+
+# Force model weights and data to float32 precision
+def cast_model_weights_to_float32(model):
+    for param in model.parameters():
+        param.data = param.data.float()  # Cast all weights to float32
+    return model
 
 # Device configuration
 device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
@@ -45,7 +44,8 @@ if torch.cuda.is_available():
 def lwm_inference(preprocessed_chs, input_type, lwm_model):
     
     dataset = prepare_for_LWM(preprocessed_chs, device)
-    
+
+    lwm_model = cast_model_weights_to_float32(lwm_model)
     # Process data through LWM
     lwm_loss, embedding_data = evaluate(lwm_model, dataset)
     print(f'LWM loss: {lwm_loss:.4f}')