Update

[ghstack-poisoned]
pytorch · XuehaiPan · Jan 7, 2025 · Jan 8, 2025 · Jan 8, 2025 · Jan 8, 2025
commit bd8994ee579054d6529e873bd0361b3f7fc8faa3
diff --git a/torch/distributed/_tools/sac_estimator.py b/torch/distributed/_tools/sac_estimator.py
@@ -429,9 +429,9 @@ def __torch_dispatch__(  # type: ignore[no-untyped-def]
         # 1. Get the runtime estimate
         out, op_time = self._estimate_runtime(func, args, kwargs)
         flat_outs = tree_iter(out)
-        out_storages_cuda: Set[UntypedStorage] = set()
-        out_storages_cpu: Set[UntypedStorage] = set()
-        cuda_devices: Set[torch.device] = set()
+        out_storages_cuda: set[UntypedStorage] = set()
+        out_storages_cpu: set[UntypedStorage] = set()
+        cuda_devices: set[torch.device] = set()
         for o in flat_outs:
             if isinstance(o, torch.Tensor):
                 if o.device.type == "cuda":

diff --git a/torch/distributed/tensor/_op_schema.py b/torch/distributed/tensor/_op_schema.py
@@ -2,7 +2,7 @@
 from collections.abc import Sequence
 from dataclasses import dataclass
 from functools import cached_property
-from typing import Any, Dict, List, Optional, Sequence, TYPE_CHECKING, Union
+from typing import Any, Optional, TYPE_CHECKING, Union
 
 import torch
 from torch._ops import OpOverload
@@ -446,7 +446,7 @@ class OpInfo:
     schema: OpSchema
     flat_args_schema: list[object]
     local_args: Sequence[object]
-    local_kwargs: Dict[str, object]
+    local_kwargs: dict[str, object]
     args_tree_spec: Optional["PyTreeSpec"] = None
 
     # the output sharding info