code: def world_size(self) -> int:
      return min(torch.cuda.device_count(), 2)
  
file_path: apex/apex/contrib/test/peer_memory/test_peer_halo_exchange_module.py
imports: import unittest
  import torch
  from torch.testing._internal import common_utils
  from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase
label_desc: Use image node_id 1 for calling the TestPeerMemory obj's underlying member method code with example usage: obj.world_size() and returns: min