code: def _test_distrib_integration(device, tol=1e-6):
      from ignite.engine import Engine
  
      rank = idist.get_rank()
      torch.manual_seed(12 + rank)
  
      def _test(metric_device):
          n_iters = 100
          batch_size = 10
  
          y_true = torch.arange(
              0, n_iters * batch_size, dtype=torch.float
          ).to(device)
          y_preds = torch.ones(
              n_iters * batch_size, dtype=torch.float
          ).to(device)
  
          def update(engine, i):
              return (
                  y_preds[i * batch_size : (i + 1) * batch_size],
                  y_true[i * batch_size : (i + 1) * batch_size],
              )
  
          engine = Engine(update)
  
          m = MeanSquaredError(device=metric_device)
          m.attach(engine, "mse")
  
          data = list(range(n_iters))
          engine.run(data=data, max_epochs=1)
  
          y_preds = idist.all_gather(y_preds)
          y_true = idist.all_gather(y_true)
  
          assert "mse" in engine.state.metrics
          res = engine.state.metrics["mse"]
  
          true_res = np.mean(
              np.power((y_true - y_preds).cpu().numpy(), 2.0)
          )
  
          assert pytest.approx(res, rel=tol) == true_res
  
      _test("cpu")
      if device.type != "xla":
          _test(idist.device())
  
imports: import os
  import numpy
  import pytest
  import torch
  import ignite.distributed
  from ignite.exceptions import NotComputableError
  from ignite.metrics import MeanSquaredError
label_desc: Use image node_id 4 for calling a global function with example usage: _test_distrib_integration(device, tol) and returns: y_preds,  y_true