no_cuda does not take effect in non distributed environment (#23795)

Signed-off-by: Wang, Yi <yi.a.wang@intel.com>
huggingface · May 26, 2023 · b7b729b · b7b729b
1 parent d61d747
commit b7b729b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/transformers/training_args.py b/src/transformers/training_args.py
@@ -1684,7 +1684,9 @@ def _setup_devices(self) -> "torch.device":
                         )
                     device = torch.device("mps")
                     self._n_gpu = 1
-
+            elif self.no_cuda:
+                device = torch.device("cpu")
+                self._n_gpu = 0
             else:
                 # if n_gpu is > 1 we'll use nn.DataParallel.
                 # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`