check world size and torch.cuda.device_count() in unit tests to try t…

…o fix contbuild (pytorch#1919)

Summary: Pull Request resolved: pytorch#1919

Differential Revision: D56448451

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)