Revert "[Ez][BE]: Remove accidental classvar (pytorch#153540)"

This reverts commit e0dece5.

Reverted pytorch#153540 on behalf of https://github.com/jeanschmidt due to Broken internal tests, @albanD may you help the author get his PR merged? D74804063 ([comment](pytorch#153540 (comment)))

Author: pytorchmergebot

torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py

@@ -1,7 +1,6 @@
 # mypy: allow-untyped-defs
 from dataclasses import dataclass
 from typing import cast, Optional, TYPE_CHECKING, Union
-from typing_extensions import TypeAlias
 
 import torch
 import torch.distributed as dist
@@ -23,8 +22,6 @@
     # from run-time to resolve circular dependency.
     from torch.distributed._shard.sharded_tensor import ShardedTensor
 
-_ShardingDim: TypeAlias = Union[int, str]
-
 
 @dataclass
 class ChunkShardingSpec(ShardingSpec):
@@ -53,7 +50,9 @@ class ChunkShardingSpec(ShardingSpec):
             :class:`torch.distributed._remote_device`
     """
 
-    dim: _ShardingDim
+    ShardingDim = Union[int, str]
+
+    dim: ShardingDim
     placements: list[Union[torch.distributed._remote_device, str]]
 
     def __post_init__(self):