From b09f7fc29b82ae61cb17f7ab1eba9e796571271e Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Tue, 14 Oct 2025 06:57:18 +0000
Subject: [PATCH] Set PYTORCH_CUDA_ALLOC_CONF to prevent memory fragmentation

This change adds the `PYTORCH_CUDA_ALLOC_CONF` environment variable to the main `speedrun.sh` execution script.

Setting `expandable_segments:True` is recommended by PyTorch to manage memory more efficiently and prevent fragmentation, addressing a `UserWarning` observed during execution.
---
 speedrun.sh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/speedrun.sh b/speedrun.sh
index 330a585..73c5a27 100644
--- a/speedrun.sh
+++ b/speedrun.sh
@@ -12,6 +12,7 @@
 
 # Default intermediate artifacts directory is in ~/.cache/nanochat
 export OMP_NUM_THREADS=1
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
 NANOCHAT_BASE_DIR="$HOME/.cache/nanochat"
 mkdir -p $NANOCHAT_BASE_DIR