diff --git a/scripts/augusta/beaker/peteish7-medlr.sh b/scripts/augusta/beaker/peteish7-medlr.sh
index 181155f3a..11a56a15a 100755
--- a/scripts/augusta/beaker/peteish7-medlr.sh
+++ b/scripts/augusta/beaker/peteish7-medlr.sh
@@ -69,7 +69,7 @@ torchrun \
       --save_interval_ephemeral=1000 \
       --eval_interval=1000 \
       --fsdp.sharding_strategy=HYBRID_SHARD \
-      --fsdp.hybrid_sharding_num_model_replicas=$NUM_NODES \
+      --fsdp.hybrid_sharding_num_model_replicas=$BEAKER_REPLICA_COUNT \
       --save_folder=$SAVE_FOLDER \
       --remote_save_folder="gs://ai2-llm/checkpoints/OLMo-medium/$NAME/" \
       --save_overwrite \