class_docstr: Specific tests on SimCLR going further than just doing a forward pass
class_docstr_tok: ["Specific","tests","on","SimCLR","going","further","than","just","doing","a","forward","pass"]
code: def worker_fn(
      gpu_id: int, world_size: int, batch_size: int, sync_file: str
  ):
      init_distributed_on_file(
          world_size=world_size, gpu_id=gpu_id, sync_file=sync_file
      )
      embeddings = torch.full(
          size=(batch_size, 3),
          fill_value=float(gpu_id),
          requires_grad=True,
      ).cuda(gpu_id)
      gathered = SimclrInfoNCECriterion.gather_embeddings(embeddings)
      if world_size == 1:
          assert gathered.equal(
              torch.tensor(
                  [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
                  device=f"cuda:{gpu_id}",
              )
          )
      if world_size == 2:
          assert gathered.equal(
              torch.tensor(
                  [
                      [0.0, 0.0, 0.0],
                      [0.0, 0.0, 0.0],
                      [1.0, 1.0, 1.0],
                      [1.0, 1.0, 1.0],
                  ],
                  device=f"cuda:{gpu_id}",
              )
          )
      assert gathered.requires_grad
  
file_path: vissl/tests/test_losses_gpu.py
imports: import unittest
  import torch
  import torch.multiprocessing
  from vissl.losses.barlow_twins_loss import BarlowTwinsCriterion
  from vissl.losses.simclr_info_nce_loss import SimclrInfoNCECriterion
  from vissl.utils.test_utils import gpu_test, init_distributed_on_file, with_temp_files
label_desc: Use image node_id 1 for calling the TestSimClrCriterionOnGpu obj's underlying member method code with example usage: obj.worker_fn(gpu_id, world_size, batch_size, sync_file) without return types