From 476be08c4aa96f8c1cae4200d2677bbe8f12cf80 Mon Sep 17 00:00:00 2001
From: Zach Mueller <muellerzr@gmail.com>
Date: Tue, 18 Jul 2023 18:40:37 -0400
Subject: [PATCH] Check for accelerate env var when doing CPU only (#24890)

Check for use-cpu
---
 src/transformers/training_args.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py
index 17adfc3355..12f0496b53 100644
--- a/src/transformers/training_args.py
+++ b/src/transformers/training_args.py
@@ -1720,7 +1720,7 @@ class TrainingArguments:
         self.distributed_state = None
         if not self.use_ipex and "ACCELERATE_USE_IPEX" not in os.environ:
             os.environ["ACCELERATE_USE_IPEX"] = "false"
-        if self.use_cpu:
+        if self.use_cpu or os.environ.get("ACCELERATE_USE_CPU", False):
             self.distributed_state = PartialState(cpu=True, backend=self.ddp_backend)
             self._n_gpu = 0
         elif is_sagemaker_mp_enabled():