disable PYTORCH_CUDA_ALLOC_CONF to see if that fixes cicd

axolotl-ai-cloud · Jul 16, 2024 · 599fbe1 · 599fbe1
1 parent 45b247c
commit 599fbe1
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/src/axolotl/train.py b/src/axolotl/train.py
@@ -53,11 +53,11 @@ def train(
     *, cfg: DictDefault, cli_args: TrainerCliArgs, dataset_meta: TrainDatasetMeta
 ) -> Tuple[Union[PeftModel, PreTrainedModel], PreTrainedTokenizer]:
     # enable expandable segments for cuda allocation to improve VRAM usage
-    torch_version = torch.__version__.split(".")
-    torch_major, torch_minor = int(torch_version[0]), int(torch_version[1])
-    if torch_major == 2 and torch_minor >= 2:
-        if os.getenv("PYTORCH_CUDA_ALLOC_CONF") is None:
-            os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+    # torch_version = torch.__version__.split(".")
+    # torch_major, torch_minor = int(torch_version[0]), int(torch_version[1])
+    # if torch_major == 2 and torch_minor >= 2:
+    #     if os.getenv("PYTORCH_CUDA_ALLOC_CONF") is None:
+    #         os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 
     # load the tokenizer first
     LOG.debug(