We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Distributed model initialisation - CIFAR 10 example
## ❓ Questions/Help/Support

In the cifar10 example when a [model](https://github.com/pytorch/ignite/blob/master/examples/contrib/cifar10/main.py#L51) is defined inside the process, is the model updated and shared between the `nproc_per_node`?

In most of the TPU examples I have seen they define the model before the `Parallel` process.
- [tpu-training-super-fast-xlmroberta](https://www.kaggle.com/shonenkov/tpu-training-super-fast-xlmroberta)
- [i-like-clean-tpu-training-kernels-i-can-not-lie](https://www.kaggle.com/abhishek/i-like-clean-tpu-training-kernels-i-can-not-lie)

Does the way ignite handles the model ensure the same model is used throughout the training?

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `ignite/distributed/auto.py`
Content:
```
1 import warnings
2 
3 import torch
4 import torch.nn as nn
5 from torch.optim.optimizer import Optimizer
6 from torch.utils.data import DataLoader, Dataset
7 from torch.utils.data.distributed import DistributedSampler
8 from torch.utils.data.sampler import Sampler
9 
10 from ignite.distributed import utils as idist
11 from ignite.distributed.comp_models import native as idist_native
12 from ignite.distributed.comp_models import xla as idist_xla
13 from ignite.utils import setup_logger
14 
15 __all__ = ["auto_dataloader", "auto_model", "auto_optim", "DistributedProxySampler"]
16 
17 
18 def auto_dataloader(dataset, **kwargs):
19     """Helper method to create a dataloader adapted for non-distributed and distributed configurations (supporting
20     all available backends from :meth:`~ignite.distributed.utils.available_backends()`).
21 
22     Internally, we create a dataloader with provided kwargs while applying the following updates:
23 
24     - batch size is scaled by world size: ``batch_size / world_size`` if larger or equal world size.
25     - number of workers is scaled by number of local processes: ``num_workers / nprocs`` if larger or equal world size.
26     - if no sampler provided by user, `torch DistributedSampler` is setup.
27     - if a sampler is provided by user, it is wrapped by :class:`~ignite.distributed.auto.DistributedProxySampler`.
28     - if the default device is 'cuda', `pin_memory` is automatically set to `True`.
29 
30     .. warning::
31 
32         Custom batch sampler is not adapted for distributed configuration. Please, make sure that provided batch
33         sampler is compatible with distributed configuration.
34 
35     Examples:
36 
37     .. code-block:: python
38 
39         import ignite.distribted as idist
40 
41         train_loader = idist.auto_dataloader(
42             train_dataset,
43             batch_size=32,
44             num_workers=4,
45             shuffle=True,
46             pin_memory="cuda" in idist.device().type,
47             drop_last=True,
48         )
49 
50     Args:
51         dataset (Dataset): input torch dataset
52         **kwargs: keyword arguments for `torch DataLoader`_.
53 
54     Returns:
55         `torch DataLoader`_ or `XLA MpDeviceLoader`_ for XLA devices
56 
57     .. _torch DataLoader: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader
58     .. _XLA MpDeviceLoader: https://github.com/pytorch/xla/blob/master/torch_xla/distributed/parallel_loader.py#L178
59     .. _torch DistributedSampler:
60         https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler
61     """
62     rank = idist.get_rank()
63     world_size = idist.get_world_size()
64 
65     logger = setup_logger(__name__ + ".auto_dataloader")
66     if world_size > 1:
67         if "batch_size" in kwargs and kwargs["batch_size"] >= world_size:
68             kwargs["batch_size"] //= world_size
69 
70         nproc = idist.get_nproc_per_node()
71         if "num_workers" in kwargs and kwargs["num_workers"] >= nproc:
72             kwargs["num_workers"] = (kwargs["num_workers"] + nproc - 1) // nproc
73 
74         if "batch_sampler" not in kwargs:
75             if kwargs.get("sampler", None) is not None:
76                 sampler = DistributedProxySampler(kwargs["sampler"], num_replicas=world_size, rank=rank)
77             else:
78                 sampler = DistributedSampler(
79                     dataset, num_replicas=world_size, rank=rank, shuffle=kwargs.get("shuffle", True)
80                 )
81                 # we need to remove "shuffle" from kwargs if sampler is used
82                 if "shuffle" in kwargs:
83                     del kwargs["shuffle"]
84 
85             kwargs["sampler"] = sampler
86         else:
87             warnings.warn(
88                 "Found batch_sampler in provided kwargs. Please, make sure that it is compatible "
89                 "with distributed configuration"
90             )
91 
92     if idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU and kwargs.get("pin_memory", False):
93         # TODO: How about XLA GPU ?
94         warnings.warn(
95             "Found incompatible options: xla support and pin_memory args equal True. "
96             "Argument `pin_memory=False` will be used to construct data loader."
97         )
98         kwargs["pin_memory"] = False
99     else:
100         kwargs["pin_memory"] = kwargs.get("pin_memory", "cuda" in idist.device().type)
101 
102     logger.info("Use data loader kwargs for dataset '{}': \n\t{}".format(repr(dataset)[:20].strip(), kwargs))
103     dataloader = DataLoader(dataset, **kwargs)
104 
105     if idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU and world_size > 1:
106 
107         logger.info("DataLoader is wrapped by `MpDeviceLoader` on XLA")
108 
109         mp_device_loader_cls = _MpDeviceLoader
110         try:
111             from torch_xla.distributed.parallel_loader import MpDeviceLoader
112 
113             mp_device_loader_cls = MpDeviceLoader
114         except ImportError:
115             pass
116 
117         sampler = dataloader.sampler
118         dataloader = mp_device_loader_cls(dataloader, idist.device())
119         dataloader.sampler = sampler
120 
121     return dataloader
122 
123 
124 def auto_model(model: nn.Module) -> nn.Module:
125     """Helper method to adapt provided model for non-distributed and distributed configurations (supporting
126     all available backends from :meth:`~ignite.distributed.utils.available_backends()`).
127 
128     Internally, we perform to following:
129 
130     - send model to current :meth:`~ignite.distributed.utils.device()`.
131     - wrap the model to `torch DistributedDataParallel`_ for native torch distributed if world size is larger than 1
132     - wrap the model to `torch DataParallel`_ if no distributed context found and more than one CUDA devices available.
133 
134     Examples:
135 
136     .. code-block:: python
137 
138         import ignite.distribted as idist
139 
140         model = idist.auto_model(model)
141 
142     Args:
143         model (torch.nn.Module): model to adapt.
144 
145     Returns:
146         torch.nn.Module
147 
148     .. _torch DistributedDataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.parallel.DistributedDataParallel
149     .. _torch DataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.DataParallel
150     """
151     logger = setup_logger(__name__ + ".auto_model")
152 
153     model.to(idist.device())
154 
155     # distributed data parallel model
156     if idist.get_world_size() > 1:
157         if idist.backend() == idist_native.NCCL:
158             lrank = idist.get_local_rank()
159             logger.info("Apply torch DistributedDataParallel on model, device id: {}".format(lrank))
160             model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[lrank,])
161         elif idist.backend() == idist_native.GLOO:
162             logger.info("Apply torch DistributedDataParallel on model")
163             model = torch.nn.parallel.DistributedDataParallel(model)
164 
165     # not distributed but multiple GPUs reachable so data parallel model
166     elif torch.cuda.device_count() > 1 and "cuda" in idist.device().type:
167         logger.info("Apply torch DataParallel on model")
168         model = torch.nn.parallel.DataParallel(model)
169 
170     return model
171 
172 
173 def auto_optim(optimizer: Optimizer) -> Optimizer:
174     """Helper method to adapt optimizer for non-distributed and distributed configurations (supporting
175     all available backends from :meth:`~ignite.distributed.utils.available_backends()`).
176 
177     Internally, this method is no-op for non-distributed and torch native distributed configuration.
178     For XLA distributed configuration, we create a new class that inherits from provided optimizer.
179     The goal is to override the `step()` method with specific `xm.optimizer_step`_ implementation.
180 
181     Examples:
182 
183     .. code-block:: python
184 
185         import ignite.distribted as idist
186 
187         optimizer = idist.auto_optim(optimizer)
188 
189 
190     Args:
191         optimizer (Optimizer): input torch optimizer
192 
193     Returns:
194         Optimizer
195 
196     .. _xm.optimizer_step: http://pytorch.org/xla/release/1.5/index.html#torch_xla.core.xla_model.optimizer_step
197 
198     """
199     if not (idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU):
200         return optimizer
201 
202     cls = type(optimizer.__class__.__name__, (optimizer.__class__,), dict(_XLADistributedOptimizer.__dict__))
203     return cls(optimizer)
204 
205 
206 class DistributedProxySampler(DistributedSampler):
207     """Distributed sampler proxy to adapt user's sampler for distributed data parallelism configuration.
208 
209     Code is based on https://github.com/pytorch/pytorch/issues/23430#issuecomment-562350407
210 
211 
212     .. note::
213         Input sampler is assumed to have a constant size.
214 
215     Args:
216         sampler (Sampler): Input torch data sampler.
217         num_replicas (int, optional): Number of processes participating in distributed training.
218         rank (int, optional): Rank of the current process within ``num_replicas``.
219 
220     """
221 
222     def __init__(self, sampler: Sampler, num_replicas=None, rank=None):
223 
224         if not isinstance(sampler, Sampler):
225             raise TypeError("Argument sampler should be instance of torch Sampler, but given: {}".format(type(sampler)))
226 
227         if not hasattr(sampler, "__len__"):
228             raise TypeError("Argument sampler should have length")
229 
230         super(DistributedProxySampler, self).__init__(sampler, num_replicas=num_replicas, rank=rank, shuffle=False)
231         self.sampler = sampler
232 
233     def __iter__(self):
234         # deterministically shuffle based on epoch
235         torch.manual_seed(self.epoch)
236 
237         indices = []
238         while len(indices) < self.total_size:
239             indices += list(self.sampler)
240 
241         if len(indices) > self.total_size:
242             indices = indices[: self.total_size]
243 
244         # subsample
245         indices = indices[self.rank : self.total_size : self.num_replicas]
246         if len(indices) != self.num_samples:
247             raise RuntimeError("{} vs {}".format(len(indices), self.num_samples))
248 
249         return iter(indices)
250 
251 
252 if idist.has_xla_support:
253 
254     import torch_xla.core.xla_model as xm
255     from torch_xla.distributed.parallel_loader import ParallelLoader
256 
257     class _MpDeviceLoader:
258         # https://github.com/pytorch/xla/pull/2117
259         # From pytorch/xla if `torch_xla.distributed.parallel_loader.MpDeviceLoader` is not available
260         def __init__(self, loader, device, **kwargs):
261             self._loader = loader
262             self._device = device
263             self._parallel_loader_kwargs = kwargs
264 
265         def __iter__(self):
266             parallel_loader = ParallelLoader(self._loader, [self._device], **self._parallel_loader_kwargs)
267             return parallel_loader.per_device_loader(self._device)
268 
269         def __len__(self):
270             return len(self._loader)
271 
272     class _XLADistributedOptimizer(Optimizer):
273         def __init__(self, optimizer):
274             super(self.__class__, self).__init__(optimizer.param_groups)
275             self.wrapped_optimizer = optimizer
276 
277         def step(self, closure=None):
278             xm.optimizer_step(self.wrapped_optimizer, barrier=True)
279
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/ignite/distributed/auto.py b/ignite/distributed/auto.py
--- a/ignite/distributed/auto.py
+++ b/ignite/distributed/auto.py
@@ -127,8 +127,8 @@
 
     Internally, we perform to following:
 
-    - send model to current :meth:`~ignite.distributed.utils.device()`.
-    - wrap the model to `torch DistributedDataParallel`_ for native torch distributed if world size is larger than 1
+    - send model to current :meth:`~ignite.distributed.utils.device()` if model's parameters are not on the device.
+    - wrap the model to `torch DistributedDataParallel`_ for native torch distributed if world size is larger than 1.
     - wrap the model to `torch DataParallel`_ if no distributed context found and more than one CUDA devices available.
 
     Examples:
@@ -139,6 +139,15 @@
 
         model = idist.auto_model(model)
 
+    In addition with NVidia/Apex, it can be used in the following way:
+
+    .. code-block:: python
+
+        import ignite.distribted as idist
+
+        model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
+        model = idist.auto_model(model)
+
     Args:
         model (torch.nn.Module): model to adapt.
 
@@ -150,7 +159,10 @@
     """
     logger = setup_logger(__name__ + ".auto_model")
 
-    model.to(idist.device())
+    # Put model's parameters to device if its parameters are not on the device
+    device = idist.device()
+    if not all([p.device == device for p in model.parameters()]):
+        model.to(device)
 
     # distributed data parallel model
     if idist.get_world_size() > 1:

{"golden_diff": "diff --git a/ignite/distributed/auto.py b/ignite/distributed/auto.py\n--- a/ignite/distributed/auto.py\n+++ b/ignite/distributed/auto.py\n@@ -127,8 +127,8 @@\n \n     Internally, we perform to following:\n \n-    - send model to current :meth:`~ignite.distributed.utils.device()`.\n-    - wrap the model to `torch DistributedDataParallel`_ for native torch distributed if world size is larger than 1\n+    - send model to current :meth:`~ignite.distributed.utils.device()` if model's parameters are not on the device.\n+    - wrap the model to `torch DistributedDataParallel`_ for native torch distributed if world size is larger than 1.\n     - wrap the model to `torch DataParallel`_ if no distributed context found and more than one CUDA devices available.\n \n     Examples:\n@@ -139,6 +139,15 @@\n \n         model = idist.auto_model(model)\n \n+    In addition with NVidia/Apex, it can be used in the following way:\n+\n+    .. code-block:: python\n+\n+        import ignite.distribted as idist\n+\n+        model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)\n+        model = idist.auto_model(model)\n+\n     Args:\n         model (torch.nn.Module): model to adapt.\n \n@@ -150,7 +159,10 @@\n     \"\"\"\n     logger = setup_logger(__name__ + \".auto_model\")\n \n-    model.to(idist.device())\n+    # Put model's parameters to device if its parameters are not on the device\n+    device = idist.device()\n+    if not all([p.device == device for p in model.parameters()]):\n+        model.to(device)\n \n     # distributed data parallel model\n     if idist.get_world_size() > 1:\n", "issue": "Distributed model initialisation - CIFAR 10 example\n## \u2753 Questions/Help/Support\r\n\r\nIn the cifar10 example when a [model](https://github.com/pytorch/ignite/blob/master/examples/contrib/cifar10/main.py#L51) is defined inside the process, is the model updated and shared between the `nproc_per_node`?\r\n\r\nIn most of the TPU examples I have seen they define the model before the `Parallel` process.\r\n- [tpu-training-super-fast-xlmroberta](https://www.kaggle.com/shonenkov/tpu-training-super-fast-xlmroberta)\r\n- [i-like-clean-tpu-training-kernels-i-can-not-lie](https://www.kaggle.com/abhishek/i-like-clean-tpu-training-kernels-i-can-not-lie)\r\n\r\nDoes the way ignite handles the model ensure the same model is used throughout the training?\n", "before_files": [{"content": "import warnings\n\nimport torch\nimport torch.nn as nn\nfrom torch.optim.optimizer import Optimizer\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch.utils.data.sampler import Sampler\n\nfrom ignite.distributed import utils as idist\nfrom ignite.distributed.comp_models import native as idist_native\nfrom ignite.distributed.comp_models import xla as idist_xla\nfrom ignite.utils import setup_logger\n\n__all__ = [\"auto_dataloader\", \"auto_model\", \"auto_optim\", \"DistributedProxySampler\"]\n\n\ndef auto_dataloader(dataset, **kwargs):\n    \"\"\"Helper method to create a dataloader adapted for non-distributed and distributed configurations (supporting\n    all available backends from :meth:`~ignite.distributed.utils.available_backends()`).\n\n    Internally, we create a dataloader with provided kwargs while applying the following updates:\n\n    - batch size is scaled by world size: ``batch_size / world_size`` if larger or equal world size.\n    - number of workers is scaled by number of local processes: ``num_workers / nprocs`` if larger or equal world size.\n    - if no sampler provided by user, `torch DistributedSampler` is setup.\n    - if a sampler is provided by user, it is wrapped by :class:`~ignite.distributed.auto.DistributedProxySampler`.\n    - if the default device is 'cuda', `pin_memory` is automatically set to `True`.\n\n    .. warning::\n\n        Custom batch sampler is not adapted for distributed configuration. Please, make sure that provided batch\n        sampler is compatible with distributed configuration.\n\n    Examples:\n\n    .. code-block:: python\n\n        import ignite.distribted as idist\n\n        train_loader = idist.auto_dataloader(\n            train_dataset,\n            batch_size=32,\n            num_workers=4,\n            shuffle=True,\n            pin_memory=\"cuda\" in idist.device().type,\n            drop_last=True,\n        )\n\n    Args:\n        dataset (Dataset): input torch dataset\n        **kwargs: keyword arguments for `torch DataLoader`_.\n\n    Returns:\n        `torch DataLoader`_ or `XLA MpDeviceLoader`_ for XLA devices\n\n    .. _torch DataLoader: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader\n    .. _XLA MpDeviceLoader: https://github.com/pytorch/xla/blob/master/torch_xla/distributed/parallel_loader.py#L178\n    .. _torch DistributedSampler:\n        https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler\n    \"\"\"\n    rank = idist.get_rank()\n    world_size = idist.get_world_size()\n\n    logger = setup_logger(__name__ + \".auto_dataloader\")\n    if world_size > 1:\n        if \"batch_size\" in kwargs and kwargs[\"batch_size\"] >= world_size:\n            kwargs[\"batch_size\"] //= world_size\n\n        nproc = idist.get_nproc_per_node()\n        if \"num_workers\" in kwargs and kwargs[\"num_workers\"] >= nproc:\n            kwargs[\"num_workers\"] = (kwargs[\"num_workers\"] + nproc - 1) // nproc\n\n        if \"batch_sampler\" not in kwargs:\n            if kwargs.get(\"sampler\", None) is not None:\n                sampler = DistributedProxySampler(kwargs[\"sampler\"], num_replicas=world_size, rank=rank)\n            else:\n                sampler = DistributedSampler(\n                    dataset, num_replicas=world_size, rank=rank, shuffle=kwargs.get(\"shuffle\", True)\n                )\n                # we need to remove \"shuffle\" from kwargs if sampler is used\n                if \"shuffle\" in kwargs:\n                    del kwargs[\"shuffle\"]\n\n            kwargs[\"sampler\"] = sampler\n        else:\n            warnings.warn(\n                \"Found batch_sampler in provided kwargs. Please, make sure that it is compatible \"\n                \"with distributed configuration\"\n            )\n\n    if idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU and kwargs.get(\"pin_memory\", False):\n        # TODO: How about XLA GPU ?\n        warnings.warn(\n            \"Found incompatible options: xla support and pin_memory args equal True. \"\n            \"Argument `pin_memory=False` will be used to construct data loader.\"\n        )\n        kwargs[\"pin_memory\"] = False\n    else:\n        kwargs[\"pin_memory\"] = kwargs.get(\"pin_memory\", \"cuda\" in idist.device().type)\n\n    logger.info(\"Use data loader kwargs for dataset '{}': \\n\\t{}\".format(repr(dataset)[:20].strip(), kwargs))\n    dataloader = DataLoader(dataset, **kwargs)\n\n    if idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU and world_size > 1:\n\n        logger.info(\"DataLoader is wrapped by `MpDeviceLoader` on XLA\")\n\n        mp_device_loader_cls = _MpDeviceLoader\n        try:\n            from torch_xla.distributed.parallel_loader import MpDeviceLoader\n\n            mp_device_loader_cls = MpDeviceLoader\n        except ImportError:\n            pass\n\n        sampler = dataloader.sampler\n        dataloader = mp_device_loader_cls(dataloader, idist.device())\n        dataloader.sampler = sampler\n\n    return dataloader\n\n\ndef auto_model(model: nn.Module) -> nn.Module:\n    \"\"\"Helper method to adapt provided model for non-distributed and distributed configurations (supporting\n    all available backends from :meth:`~ignite.distributed.utils.available_backends()`).\n\n    Internally, we perform to following:\n\n    - send model to current :meth:`~ignite.distributed.utils.device()`.\n    - wrap the model to `torch DistributedDataParallel`_ for native torch distributed if world size is larger than 1\n    - wrap the model to `torch DataParallel`_ if no distributed context found and more than one CUDA devices available.\n\n    Examples:\n\n    .. code-block:: python\n\n        import ignite.distribted as idist\n\n        model = idist.auto_model(model)\n\n    Args:\n        model (torch.nn.Module): model to adapt.\n\n    Returns:\n        torch.nn.Module\n\n    .. _torch DistributedDataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.parallel.DistributedDataParallel\n    .. _torch DataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.DataParallel\n    \"\"\"\n    logger = setup_logger(__name__ + \".auto_model\")\n\n    model.to(idist.device())\n\n    # distributed data parallel model\n    if idist.get_world_size() > 1:\n        if idist.backend() == idist_native.NCCL:\n            lrank = idist.get_local_rank()\n            logger.info(\"Apply torch DistributedDataParallel on model, device id: {}\".format(lrank))\n            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[lrank,])\n        elif idist.backend() == idist_native.GLOO:\n            logger.info(\"Apply torch DistributedDataParallel on model\")\n            model = torch.nn.parallel.DistributedDataParallel(model)\n\n    # not distributed but multiple GPUs reachable so data parallel model\n    elif torch.cuda.device_count() > 1 and \"cuda\" in idist.device().type:\n        logger.info(\"Apply torch DataParallel on model\")\n        model = torch.nn.parallel.DataParallel(model)\n\n    return model\n\n\ndef auto_optim(optimizer: Optimizer) -> Optimizer:\n    \"\"\"Helper method to adapt optimizer for non-distributed and distributed configurations (supporting\n    all available backends from :meth:`~ignite.distributed.utils.available_backends()`).\n\n    Internally, this method is no-op for non-distributed and torch native distributed configuration.\n    For XLA distributed configuration, we create a new class that inherits from provided optimizer.\n    The goal is to override the `step()` method with specific `xm.optimizer_step`_ implementation.\n\n    Examples:\n\n    .. code-block:: python\n\n        import ignite.distribted as idist\n\n        optimizer = idist.auto_optim(optimizer)\n\n\n    Args:\n        optimizer (Optimizer): input torch optimizer\n\n    Returns:\n        Optimizer\n\n    .. _xm.optimizer_step: http://pytorch.org/xla/release/1.5/index.html#torch_xla.core.xla_model.optimizer_step\n\n    \"\"\"\n    if not (idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU):\n        return optimizer\n\n    cls = type(optimizer.__class__.__name__, (optimizer.__class__,), dict(_XLADistributedOptimizer.__dict__))\n    return cls(optimizer)\n\n\nclass DistributedProxySampler(DistributedSampler):\n    \"\"\"Distributed sampler proxy to adapt user's sampler for distributed data parallelism configuration.\n\n    Code is based on https://github.com/pytorch/pytorch/issues/23430#issuecomment-562350407\n\n\n    .. note::\n        Input sampler is assumed to have a constant size.\n\n    Args:\n        sampler (Sampler): Input torch data sampler.\n        num_replicas (int, optional): Number of processes participating in distributed training.\n        rank (int, optional): Rank of the current process within ``num_replicas``.\n\n    \"\"\"\n\n    def __init__(self, sampler: Sampler, num_replicas=None, rank=None):\n\n        if not isinstance(sampler, Sampler):\n            raise TypeError(\"Argument sampler should be instance of torch Sampler, but given: {}\".format(type(sampler)))\n\n        if not hasattr(sampler, \"__len__\"):\n            raise TypeError(\"Argument sampler should have length\")\n\n        super(DistributedProxySampler, self).__init__(sampler, num_replicas=num_replicas, rank=rank, shuffle=False)\n        self.sampler = sampler\n\n    def __iter__(self):\n        # deterministically shuffle based on epoch\n        torch.manual_seed(self.epoch)\n\n        indices = []\n        while len(indices) < self.total_size:\n            indices += list(self.sampler)\n\n        if len(indices) > self.total_size:\n            indices = indices[: self.total_size]\n\n        # subsample\n        indices = indices[self.rank : self.total_size : self.num_replicas]\n        if len(indices) != self.num_samples:\n            raise RuntimeError(\"{} vs {}\".format(len(indices), self.num_samples))\n\n        return iter(indices)\n\n\nif idist.has_xla_support:\n\n    import torch_xla.core.xla_model as xm\n    from torch_xla.distributed.parallel_loader import ParallelLoader\n\n    class _MpDeviceLoader:\n        # https://github.com/pytorch/xla/pull/2117\n        # From pytorch/xla if `torch_xla.distributed.parallel_loader.MpDeviceLoader` is not available\n        def __init__(self, loader, device, **kwargs):\n            self._loader = loader\n            self._device = device\n            self._parallel_loader_kwargs = kwargs\n\n        def __iter__(self):\n            parallel_loader = ParallelLoader(self._loader, [self._device], **self._parallel_loader_kwargs)\n            return parallel_loader.per_device_loader(self._device)\n\n        def __len__(self):\n            return len(self._loader)\n\n    class _XLADistributedOptimizer(Optimizer):\n        def __init__(self, optimizer):\n            super(self.__class__, self).__init__(optimizer.param_groups)\n            self.wrapped_optimizer = optimizer\n\n        def step(self, closure=None):\n            xm.optimizer_step(self.wrapped_optimizer, barrier=True)\n", "path": "ignite/distributed/auto.py"}], "after_files": [{"content": "import warnings\n\nimport torch\nimport torch.nn as nn\nfrom torch.optim.optimizer import Optimizer\nfrom torch.utils.data import DataLoader, Dataset\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch.utils.data.sampler import Sampler\n\nfrom ignite.distributed import utils as idist\nfrom ignite.distributed.comp_models import native as idist_native\nfrom ignite.distributed.comp_models import xla as idist_xla\nfrom ignite.utils import setup_logger\n\n__all__ = [\"auto_dataloader\", \"auto_model\", \"auto_optim\", \"DistributedProxySampler\"]\n\n\ndef auto_dataloader(dataset, **kwargs):\n    \"\"\"Helper method to create a dataloader adapted for non-distributed and distributed configurations (supporting\n    all available backends from :meth:`~ignite.distributed.utils.available_backends()`).\n\n    Internally, we create a dataloader with provided kwargs while applying the following updates:\n\n    - batch size is scaled by world size: ``batch_size / world_size`` if larger or equal world size.\n    - number of workers is scaled by number of local processes: ``num_workers / nprocs`` if larger or equal world size.\n    - if no sampler provided by user, `torch DistributedSampler` is setup.\n    - if a sampler is provided by user, it is wrapped by :class:`~ignite.distributed.auto.DistributedProxySampler`.\n    - if the default device is 'cuda', `pin_memory` is automatically set to `True`.\n\n    .. warning::\n\n        Custom batch sampler is not adapted for distributed configuration. Please, make sure that provided batch\n        sampler is compatible with distributed configuration.\n\n    Examples:\n\n    .. code-block:: python\n\n        import ignite.distribted as idist\n\n        train_loader = idist.auto_dataloader(\n            train_dataset,\n            batch_size=32,\n            num_workers=4,\n            shuffle=True,\n            pin_memory=\"cuda\" in idist.device().type,\n            drop_last=True,\n        )\n\n    Args:\n        dataset (Dataset): input torch dataset\n        **kwargs: keyword arguments for `torch DataLoader`_.\n\n    Returns:\n        `torch DataLoader`_ or `XLA MpDeviceLoader`_ for XLA devices\n\n    .. _torch DataLoader: https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader\n    .. _XLA MpDeviceLoader: https://github.com/pytorch/xla/blob/master/torch_xla/distributed/parallel_loader.py#L178\n    .. _torch DistributedSampler:\n        https://pytorch.org/docs/stable/data.html#torch.utils.data.distributed.DistributedSampler\n    \"\"\"\n    rank = idist.get_rank()\n    world_size = idist.get_world_size()\n\n    logger = setup_logger(__name__ + \".auto_dataloader\")\n    if world_size > 1:\n        if \"batch_size\" in kwargs and kwargs[\"batch_size\"] >= world_size:\n            kwargs[\"batch_size\"] //= world_size\n\n        nproc = idist.get_nproc_per_node()\n        if \"num_workers\" in kwargs and kwargs[\"num_workers\"] >= nproc:\n            kwargs[\"num_workers\"] = (kwargs[\"num_workers\"] + nproc - 1) // nproc\n\n        if \"batch_sampler\" not in kwargs:\n            if kwargs.get(\"sampler\", None) is not None:\n                sampler = DistributedProxySampler(kwargs[\"sampler\"], num_replicas=world_size, rank=rank)\n            else:\n                sampler = DistributedSampler(\n                    dataset, num_replicas=world_size, rank=rank, shuffle=kwargs.get(\"shuffle\", True)\n                )\n                # we need to remove \"shuffle\" from kwargs if sampler is used\n                if \"shuffle\" in kwargs:\n                    del kwargs[\"shuffle\"]\n\n            kwargs[\"sampler\"] = sampler\n        else:\n            warnings.warn(\n                \"Found batch_sampler in provided kwargs. Please, make sure that it is compatible \"\n                \"with distributed configuration\"\n            )\n\n    if idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU and kwargs.get(\"pin_memory\", False):\n        # TODO: How about XLA GPU ?\n        warnings.warn(\n            \"Found incompatible options: xla support and pin_memory args equal True. \"\n            \"Argument `pin_memory=False` will be used to construct data loader.\"\n        )\n        kwargs[\"pin_memory\"] = False\n    else:\n        kwargs[\"pin_memory\"] = kwargs.get(\"pin_memory\", \"cuda\" in idist.device().type)\n\n    logger.info(\"Use data loader kwargs for dataset '{}': \\n\\t{}\".format(repr(dataset)[:20].strip(), kwargs))\n    dataloader = DataLoader(dataset, **kwargs)\n\n    if idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU and world_size > 1:\n\n        logger.info(\"DataLoader is wrapped by `MpDeviceLoader` on XLA\")\n\n        mp_device_loader_cls = _MpDeviceLoader\n        try:\n            from torch_xla.distributed.parallel_loader import MpDeviceLoader\n\n            mp_device_loader_cls = MpDeviceLoader\n        except ImportError:\n            pass\n\n        sampler = dataloader.sampler\n        dataloader = mp_device_loader_cls(dataloader, idist.device())\n        dataloader.sampler = sampler\n\n    return dataloader\n\n\ndef auto_model(model: nn.Module) -> nn.Module:\n    \"\"\"Helper method to adapt provided model for non-distributed and distributed configurations (supporting\n    all available backends from :meth:`~ignite.distributed.utils.available_backends()`).\n\n    Internally, we perform to following:\n\n    - send model to current :meth:`~ignite.distributed.utils.device()` if model's parameters are not on the device.\n    - wrap the model to `torch DistributedDataParallel`_ for native torch distributed if world size is larger than 1.\n    - wrap the model to `torch DataParallel`_ if no distributed context found and more than one CUDA devices available.\n\n    Examples:\n\n    .. code-block:: python\n\n        import ignite.distribted as idist\n\n        model = idist.auto_model(model)\n\n    In addition with NVidia/Apex, it can be used in the following way:\n\n    .. code-block:: python\n\n        import ignite.distribted as idist\n\n        model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)\n        model = idist.auto_model(model)\n\n    Args:\n        model (torch.nn.Module): model to adapt.\n\n    Returns:\n        torch.nn.Module\n\n    .. _torch DistributedDataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.parallel.DistributedDataParallel\n    .. _torch DataParallel: https://pytorch.org/docs/stable/nn.html#torch.nn.DataParallel\n    \"\"\"\n    logger = setup_logger(__name__ + \".auto_model\")\n\n    # Put model's parameters to device if its parameters are not on the device\n    device = idist.device()\n    if not all([p.device == device for p in model.parameters()]):\n        model.to(device)\n\n    # distributed data parallel model\n    if idist.get_world_size() > 1:\n        if idist.backend() == idist_native.NCCL:\n            lrank = idist.get_local_rank()\n            logger.info(\"Apply torch DistributedDataParallel on model, device id: {}\".format(lrank))\n            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[lrank,])\n        elif idist.backend() == idist_native.GLOO:\n            logger.info(\"Apply torch DistributedDataParallel on model\")\n            model = torch.nn.parallel.DistributedDataParallel(model)\n\n    # not distributed but multiple GPUs reachable so data parallel model\n    elif torch.cuda.device_count() > 1 and \"cuda\" in idist.device().type:\n        logger.info(\"Apply torch DataParallel on model\")\n        model = torch.nn.parallel.DataParallel(model)\n\n    return model\n\n\ndef auto_optim(optimizer: Optimizer) -> Optimizer:\n    \"\"\"Helper method to adapt optimizer for non-distributed and distributed configurations (supporting\n    all available backends from :meth:`~ignite.distributed.utils.available_backends()`).\n\n    Internally, this method is no-op for non-distributed and torch native distributed configuration.\n    For XLA distributed configuration, we create a new class that inherits from provided optimizer.\n    The goal is to override the `step()` method with specific `xm.optimizer_step`_ implementation.\n\n    Examples:\n\n    .. code-block:: python\n\n        import ignite.distribted as idist\n\n        optimizer = idist.auto_optim(optimizer)\n\n\n    Args:\n        optimizer (Optimizer): input torch optimizer\n\n    Returns:\n        Optimizer\n\n    .. _xm.optimizer_step: http://pytorch.org/xla/release/1.5/index.html#torch_xla.core.xla_model.optimizer_step\n\n    \"\"\"\n    if not (idist.has_xla_support and idist.backend() == idist_xla.XLA_TPU):\n        return optimizer\n\n    cls = type(optimizer.__class__.__name__, (optimizer.__class__,), dict(_XLADistributedOptimizer.__dict__))\n    return cls(optimizer)\n\n\nclass DistributedProxySampler(DistributedSampler):\n    \"\"\"Distributed sampler proxy to adapt user's sampler for distributed data parallelism configuration.\n\n    Code is based on https://github.com/pytorch/pytorch/issues/23430#issuecomment-562350407\n\n\n    .. note::\n        Input sampler is assumed to have a constant size.\n\n    Args:\n        sampler (Sampler): Input torch data sampler.\n        num_replicas (int, optional): Number of processes participating in distributed training.\n        rank (int, optional): Rank of the current process within ``num_replicas``.\n\n    \"\"\"\n\n    def __init__(self, sampler: Sampler, num_replicas=None, rank=None):\n\n        if not isinstance(sampler, Sampler):\n            raise TypeError(\"Argument sampler should be instance of torch Sampler, but given: {}\".format(type(sampler)))\n\n        if not hasattr(sampler, \"__len__\"):\n            raise TypeError(\"Argument sampler should have length\")\n\n        super(DistributedProxySampler, self).__init__(sampler, num_replicas=num_replicas, rank=rank, shuffle=False)\n        self.sampler = sampler\n\n    def __iter__(self):\n        # deterministically shuffle based on epoch\n        torch.manual_seed(self.epoch)\n\n        indices = []\n        while len(indices) < self.total_size:\n            indices += list(self.sampler)\n\n        if len(indices) > self.total_size:\n            indices = indices[: self.total_size]\n\n        # subsample\n        indices = indices[self.rank : self.total_size : self.num_replicas]\n        if len(indices) != self.num_samples:\n            raise RuntimeError(\"{} vs {}\".format(len(indices), self.num_samples))\n\n        return iter(indices)\n\n\nif idist.has_xla_support:\n\n    import torch_xla.core.xla_model as xm\n    from torch_xla.distributed.parallel_loader import ParallelLoader\n\n    class _MpDeviceLoader:\n        # https://github.com/pytorch/xla/pull/2117\n        # From pytorch/xla if `torch_xla.distributed.parallel_loader.MpDeviceLoader` is not available\n        def __init__(self, loader, device, **kwargs):\n            self._loader = loader\n            self._device = device\n            self._parallel_loader_kwargs = kwargs\n\n        def __iter__(self):\n            parallel_loader = ParallelLoader(self._loader, [self._device], **self._parallel_loader_kwargs)\n            return parallel_loader.per_device_loader(self._device)\n\n        def __len__(self):\n            return len(self._loader)\n\n    class _XLADistributedOptimizer(Optimizer):\n        def __init__(self, optimizer):\n            super(self.__class__, self).__init__(optimizer.param_groups)\n            self.wrapped_optimizer = optimizer\n\n        def step(self, closure=None):\n            xm.optimizer_step(self.wrapped_optimizer, barrier=True)\n", "path": "ignite/distributed/auto.py"}]}