saymrwulf/pytorch
1
0
Fork 0
mirror of https://github.com/saymrwulf/pytorch.git synced 2026-05-14 20:57:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add type for device_mesh

Browse source
This commit is contained in:
Mauricio Villegas 2024-09-30 10:03:26 +02:00
parent 45cd195f24
commit 7a79dbdc05
1 changed files with 2 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
torch/nn/parallel/distributed.py
View file

@ -56,6 +56,7 @@ if dist.rpc.is_available():
from torch.distributed.rpc import RRef
if TYPE_CHECKING:
from torch.distributed.device_mesh import DeviceMesh
from torch.utils.hooks import RemovableHandle
@ -662,7 +663,7 @@ class DistributedDataParallel(Module, Joinable):
delay_all_reduce_named_params=None,
param_to_hook_all_reduce=None,
mixed_precision: Optional[_MixedPrecision] = None,
device_mesh=None,
device_mesh: "Optional[DeviceMesh]" = None,
) -> None:
super().__init__()
Joinable.__init__(self)