restart 2gpu

mwalmsley · Nov 22, 2023 · c251988 · c251988
1 parent 52718f7
commit c251988
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/only_for_me/narval/train.sh b/only_for_me/narval/train.sh
@@ -1,10 +1,10 @@
 #!/bin/bash
 #SBATCH --time=23:30:0  
 #SBATCH --nodes=1
-#SBATCH --ntasks-per-node=1
+#SBATCH --ntasks-per-node=2
 #SBATCH --cpus-per-task=10
 #SBATCH --mem-per-cpu 4G
-#SBATCH --gres=gpu:v100:1
+#SBATCH --gres=gpu:v100:2
 
 nvidia-smi
 
@@ -21,10 +21,10 @@ export NCCL_BLOCKING_WAIT=1  #Set this environment variable if you wish to use t
 
 REPO_DIR=/project/def-bovy/walml/zoobot
 srun $PYTHON $REPO_DIR/only_for_me/narval/train.py \
-    --save-dir $REPO_DIR/only_for_me/narval/desi_f512 \
+    --save-dir $REPO_DIR/only_for_me/narval/desi_f128_2gpu \
     --batch-size 256 \
-    --num-features 512 \
-    --gpus 1 \
+    --num-features 128 \
+    --gpus 2 \
     --num-workers 10 \
     --color --wandb --mixed-precision --compile-encoder
 

diff --git a/zoobot/pytorch/training/train_with_pytorch_lightning.py b/zoobot/pytorch/training/train_with_pytorch_lightning.py
@@ -115,6 +115,7 @@ def train_default_zoobot_from_scratch(
     assert save_dir is not None
     if not os.path.isdir(save_dir):
         os.mkdir(save_dir)
+    logging.info(f'Saving to {save_dir}')
 
     if color:
         logging.warning(