Update ddp_tutorial.rst (pytorch#1618)

Relax GPU count to support devices with 2 GPUs

Co-authored-by: Brian Johnson <[email protected]>

Author: bowangbj

intermediate_source/ddp_tutorial.rst

@@ -265,8 +265,8 @@ either the application or the model ``forward()`` method.
         setup(rank, world_size)
 
         # setup mp_model and devices for this process
-        dev0 = rank * 2
-        dev1 = rank * 2 + 1
+        dev0 = (rank * 2) % world_size
+        dev1 = (rank * 2 + 1) % world_size
         mp_model = ToyMpModel(dev0, dev1)
         ddp_mp_model = DDP(mp_model)
 
@@ -285,9 +285,8 @@ either the application or the model ``forward()`` method.
 
     if __name__ == "__main__":
         n_gpus = torch.cuda.device_count()
-        if n_gpus < 8:
-          print(f"Requires at least 8 GPUs to run, but got {n_gpus}.")
-        else:
-          run_demo(demo_basic, 8)
-          run_demo(demo_checkpoint, 8)
-          run_demo(demo_model_parallel, 4)
+        assert n_gpus >= 2, f"Requires at least 2 GPUs to run, but got {n_gpus}"
+        world_size = n_gpus
+        run_demo(demo_basic, world_size)
+        run_demo(demo_checkpoint, world_size)
+        run_demo(demo_model_parallel, world_size)