Code changes

inference_brain2vec.py

@@ -119,7 +119,7 @@ class Brain2vec(AutoencoderKL):
         if checkpoint_path is not None:
             if not os.path.exists(checkpoint_path):
                 raise FileNotFoundError(f"Checkpoint {checkpoint_path} not found.")
-            state_dict = torch.load(checkpoint_path, map_location=device)
+            state_dict = torch.load(checkpoint_path, map_location=device, weights_only=False)
             model.load_state_dict(state_dict)
 
         model.to(device)