code: def test_gather_all_same_size(self):
      ts = torch.full((10, 5), DistributedHelper.rank, dtype=torch.long)
      device = DistributedHelper.make_device()
  
      if device.type == "cuda":
          # Additional test: tensors do not need to be on the default device
          DistributedHelper.gather_all(ts)
  
          # On the other hand, PyTorch all_gather requires tensors to be on
          # the default device
          with self.assertRaises(Exception):
              out_t = [
                  torch.empty_like(ts)
                  for _ in range(DistributedHelper.world_size)
              ]
              torch.distributed.all_gather(out_t, ts)
  
          # ... while this should work
          out_t = [
              torch.empty_like(ts).to(device)
              for _ in range(DistributedHelper.world_size)
          ]
          torch.distributed.all_gather(out_t, ts.to(device))
  
      ts = ts.to(device)
  
      for same_shape in [False, True]:
          print(f"same_shape={same_shape}")
          # with self.subTest(same_shape=same_shape):
          tensor_list = DistributedHelper.gather_all(
              ts, same_shape=same_shape
          )
  
          self.assertEqual(
              DistributedHelper.world_size, len(tensor_list)
          )
  
          for t in tensor_list:
              self.assertEqual(device, t.device)
  
          for rank in range(DistributedHelper.world_size):
              expect = torch.full((10, 5), rank, dtype=torch.long).to(
                  device
              )
              self.assertTrue(torch.equal(tensor_list[rank], expect))
  
imports: import os
  import random
  import shutil
  import tempfile
  import time
  import unittest
  import numpy
  import torch
  import torch.distributed
  from torch.nn import Module
  from torch.nn.parallel import DistributedDataParallel
  from avalanche.benchmarks.scenarios.deprecated.generators import dataset_benchmark
  from avalanche.benchmarks.utils.classification_dataset import _make_taskaware_tensor_classification_dataset
  from avalanche.distributed import DistributedHelper
  from avalanche.distributed.distributed_helper import RollingSeedContext, BroadcastSeedContext
  from avalanche.models import SimpleMLP, as_multitask
  from avalanche.models.utils import avalanche_model_adaptation
  from avalanche.training.determinism.rng_manager import RNGManager
  from tests.distributed.distributed_test_utils import check_skip_distributed_slow_test, check_skip_distributed_test, suppress_dst_tests_output, common_dst_tests_setup
label_desc: Use image node_id 7 for calling the DistributedHelperTests obj's underlying member method code with example usage: obj.test_gather_all_same_size() without return types