update ddp test

Author: awaelchli

tests/plugins/test_ddp_plugin.py

@@ -11,6 +11,8 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from unittest import mock
+
 import torch
 from torch.nn.parallel import DistributedDataParallel
 
@@ -61,13 +63,16 @@ def on_train_start(self):
 
 
 @RunIf(min_gpus=4, special=True)
-def test_ddp_barrier_non_consecutive_device_ids(tmpdir):
-
+@mock.patch("torch.distributed.barrier")
+def test_ddp_barrier_non_consecutive_device_ids(barrier_mock, tmpdir):
+    """ Test correct usage of barriers when device ids do not start at 0 or are not consecutive. """
     model = BoringModel()
+    gpus = [1, 3]
     trainer = Trainer(
         default_root_dir=tmpdir,
         max_steps=1,
-        gpus=[1, 3],
+        gpus=gpus,
         accelerator="ddp",
     )
     trainer.fit(model)
+    barrier_mock.assert_any_call(device_ids=[gpus[trainer.local_rank]])