Convert GPU to CPU if CUDA not available

Summary: Pull Request resolved: https://github.com/facebookresearch/d2go/pull/569

Reviewed By: wat3rBro

Differential Revision: D46498855

fbshipit-source-id: 99888f6a36a0f69155c3447cc080392ae9886539

d2go/distributed.py

@@ -56,6 +56,21 @@ def get_num_processes_per_machine() -> int:
     return mcv_comm.get_local_size()
 
 
+def _maybe_convert_to_cpu_run(args, backend):
+    if get_launch_environment() == "local" and not torch.cuda.is_available():
+        assert len(args) > 0, args
+        cfg = args[0]
+        if isinstance(cfg, CfgNode) and cfg.MODEL.DEVICE == "cuda":
+            logger.warning(
+                "Detected that CUDA is not available on this machine, set MODEL.DEVICE"
+                " to cpu and backend to GLOO"
+            )
+            with temp_defrost(cfg):
+                cfg.MODEL.DEVICE = "cpu"
+        backend = "GLOO"
+    return args, backend
+
+
 # Modify mobile_cv's `default_distributed_worker` to also setup D2's comm module
 def distributed_worker(
     main_func: Callable[..., _RT],
@@ -74,6 +89,8 @@ def distributed_worker(
         )  # set the global shared context from the args passed in by mp spawn
     dist_params = dist_params or DistributedParams.from_environ()
 
+    args, backend = _maybe_convert_to_cpu_run(args, backend)
+
     with enable_dist_process_groups(backend, init_method, dist_params, timeout):
         d2_comm._LOCAL_PROCESS_GROUP = mcv_comm._LOCAL_PROCESS_GROUP
         # Now the D2's comm module should be fully functional
@@ -107,17 +124,7 @@ def launch(
         - Automatically convert GPU to CPU if CUDA is not available.
         - Add D2Go-specific initialziation in the _distributed_worker.
     """
-    if get_launch_environment() == "local" and not torch.cuda.is_available():
-        assert len(args) > 0, args
-        cfg = args[0]
-        if isinstance(cfg, CfgNode) and cfg.MODEL.DEVICE == "cuda":
-            logger.warning(
-                "Detected that CUDA is not available on this machine, set MODEL.DEVICE"
-                " to cpu and backend to GLOO"
-            )
-            with temp_defrost(cfg):
-                cfg.MODEL.DEVICE = "cpu"
-        backend = "GLOO"
+    args, backend = _maybe_convert_to_cpu_run(args, backend)
 
     return _launch(
         main_func=main_func,