From 3525c35d44bd14ce7842450f386138ffecbd65b2 Mon Sep 17 00:00:00 2001 From: Zach Mueller Date: Tue, 18 Jul 2023 12:01:56 -0400 Subject: [PATCH] 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 3a30adbe0b07..bb0a4187f8c9 100644 --- a/src/transformers/training_args.py +++ b/src/transformers/training_args.py @@ -1718,7 +1718,7 @@ def _setup_devices(self) -> "torch.device": ) AcceleratorState._reset_state(reset_partial_state=True) self.distributed_state = None - 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():