Improve check in world size; None is valid group (#4)

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

transformer_engine/pytorch/distributed.py

@@ -58,13 +58,14 @@ def set_tensor_model_parallel_attributes(
 
 def get_distributed_world_size(group: Optional[dist_group_type] = None) -> int:
     """Return world size for the distributed group."""
-    if group is None:
+    if not torch.distributed.is_initialized():
         return 1
     return torch.distributed.get_world_size(group=group)
 
 
 def get_distributed_rank(group: Optional[dist_group_type] = None) -> int:
     """Return my rank for the distributed group."""
+    assert torch.distributed.is_initialized(), "torch.distributed is not initialized."
     return torch.distributed.get_rank(group=group)