check world size and torch.cuda.device_count() in unit tests to try to fix contbuild (#1919)

Summary: Pull Request resolved: #1919

Reviewed By: sarckk

Differential Revision: D56448451

fbshipit-source-id: 2dd47f7c26e4a329278c640f8ace084b9789127e

Author: henrylhtsang

torchrec/distributed/test_utils/multi_process.py

@@ -38,7 +38,7 @@ def __init__(
         self.backend = backend
         self.local_size = local_size
 
-        if torch.cuda.is_available():
+        if torch.cuda.is_available() and world_size <= torch.cuda.device_count():
             self.device: torch.device = torch.device(f"cuda:{rank}")
             torch.cuda.set_device(self.device)