32027652

diffusion_test.py

@@ -713,6 +713,7 @@ if __name__ == "__main__":
     parser.add_argument("--autocast", type=int, required=False, default=False)
     parser.add_argument("--use_checkpoint", type=int, required=False, default=False)
     parser.add_argument("--dropout", type=float, required=False, default=0)
+    parser.add_argument("--lrate", type=float, required=False, default=1e-4)
 
     args = parser.parse_args()
 
@@ -731,6 +732,7 @@ if __name__ == "__main__":
     config.autocast = bool(args.autocast)
     config.use_checkpoint = bool(args.use_checkpoint)
     config.dropout = args.dropout
+    config.lrate = args.lrate
 
     ############################ training ################################
     if args.train:

perlmutter_diffusion_test.sbatch

@@ -38,6 +38,7 @@ srun python diffusion_test.py \
     --autocast 1 \
     --use_checkpoint 1 \
     --dropout 0.1 \
+    --lrate 7e-5 \
     --train "$SCRATCH/LEN128-DIM64-CUB16-Tvir[4, 6]-zeta[10, 250]-0809-123640.h5" \
     #--resume ./outputs/model-N1600-device_count4-node4-epoch9-20051216 \