Code changes

README.md

@@ -44,6 +44,12 @@ python create_csv.py
 mkdir ae_cache
 mkdir ae_output
 
+# install git lfs to pull large model weights
+sudo apt-get update
+sudo apt install git-lfs
+git lfs install
+git lfs pull
+
 # train the model
 nohup python brain2vec.py train \
   --dataset_csv /home/ubuntu/brain2vec/inputs.csv \

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, weights_only=False)
+            state_dict = torch.load(checkpoint_path, map_location=device)
             model.load_state_dict(state_dict)
 
         model.to(device)