Apply suggestions from code review

Co-authored-by: Nicki Skafte Detlefsen <[email protected]>

Author: Borda

src/lightning/fabric/utilities/distributed.py

@@ -319,6 +319,7 @@ def _destroy_dist_connection() -> None:
 
 
 def _get_default_process_group_backend_for_device(device: torch.device) -> str:
+    """Return corresponding distributed backend for a given device."""
     device_backend_map = torch.distributed.Backend.default_device_backend_map
     if device.type in device_backend_map:
         return device_backend_map[device.type]

tests/tests_fabric/utilities/test_distributed.py

@@ -245,6 +245,7 @@ def test_init_dist_connection_registers_destruction_handler(_, atexit_mock):
 
 
 def test_get_default_process_group_backend_for_device():
+    """Test that each device type maps to its correct default process group backend."""
     # register a custom backend for test
     torch.utils.rename_privateuse1_backend("pcu")