Fix torch.distributed._* import statements in tests (#11416)

Co-authored-by: Carlos Mocholi <carlossmocholi@gmail.com>

tests/core/test_lightning_module.py

@@ -15,7 +15,6 @@
 
 import pytest
 import torch
-import torch.distributed as dist
 from torch import nn
 from torch.optim import Adam, SGD
 
@@ -304,16 +303,18 @@ def assert_device(device: torch.device) -> None:
     assert_device(torch.device("cpu"))
 
 
-class BoringModelWithShardedTensor(BoringModel):
-    def __init__(self, spec):
-        super().__init__()
-        self.sharded_tensor = dist._sharded_tensor.empty(spec, 10, 20)
-        self.sharded_tensor.local_shards()[0].tensor.fill_(0)
-
-
 @RunIf(min_torch="1.10", skip_windows=True)
 def test_sharded_tensor_state_dict(tmpdir, single_process_pg):
-    spec = dist._sharding_spec.ChunkShardingSpec(
+    from torch.distributed._sharded_tensor import empty as sharded_tensor_empty
+    from torch.distributed._sharding_spec import ChunkShardingSpec
+
+    class BoringModelWithShardedTensor(BoringModel):
+        def __init__(self, spec):
+            super().__init__()
+            self.sharded_tensor = sharded_tensor_empty(spec, 10, 20)
+            self.sharded_tensor.local_shards()[0].tensor.fill_(0)
+
+    spec = ChunkShardingSpec(
         dim=0,
         placements=[
             "rank:0/cpu",