code: def run_cluster_assignment(self, with_fsdp: bool):
      with in_temporary_directory() as pretrain_dir:
          # Pre-train a SwAV model in order to get some weights
          pretrain_config = self._create_pretraining_config(
              with_fsdp=with_fsdp
          )
          run_integration_test(pretrain_config)
  
          # Extract the cluster assignments of each sample
          with in_temporary_directory() as extract_dir:
              extract_config = self._create_extract_cluster_config(
                  with_fsdp=with_fsdp,
                  checkpoint_path=os.path.join(
                      pretrain_dir, "checkpoint.torch"
                  ),
              )
              extract_config.EXTRACT_FEATURES.CHUNK_THRESHOLD = 10
              run_integration_test(
                  extract_config, engine_name="extract_cluster"
              )
              extraction_outputs = os.listdir(extract_dir)
  
              # Check that the cluster assignments are computed in both
              # compact format and dataset disk_filelist format
              self.assertIn(
                  "cluster_assignments.torch", extraction_outputs
              )
              self.assertIn("train_images.npy", extraction_outputs)
              self.assertIn("train_labels.npy", extraction_outputs)
              self.assertIn("test_images.npy", extraction_outputs)
              self.assertIn("test_labels.npy", extraction_outputs)
  
              # Check that the soft assignments (on prototypes) are exported
              for rank in range(2):
                  for chunk in range(2):
                      file_name = f"rank{rank}_chunk{chunk}_train_heads_protos.npy"
                      self.assertIn(file_name, extraction_outputs)
                      self.assertEqual(
                          np.load(file_name).shape[1], 3000
                      )
                  file_name = f"rank{rank}_chunk0_test_heads_protos.npy"
                  self.assertIn(file_name, extraction_outputs)
                  self.assertEqual(np.load(file_name).shape[1], 3000)
  
              # Copy the cluster assignments
              shutil.move(
                  src=os.path.join(
                      extract_dir, "cluster_assignments.torch"
                  ),
                  dst=os.path.join(
                      pretrain_dir, "cluster_assignments.torch"
                  ),
              )
  
          # Load the cluster assignments and check their structure
          assignments = ClusterAssignmentLoader.load_cluster_assigment(
              "cluster_assignments.torch"
          )
          self.assertEqual(
              40, len(assignments.cluster_assignments["TRAIN"])
          )
          self.assertEqual(
              20, len(assignments.cluster_assignments["TEST"])
          )
  
file_path: vissl/tests/test_extract_cluster.py
imports: import os
  import shutil
  import unittest
  import numpy
  from vissl.utils.cluster_utils import ClusterAssignmentLoader
  from vissl.utils.hydra_config import compose_hydra_configuration, convert_to_attrdict
  from vissl.utils.test_utils import gpu_test, in_temporary_directory, run_integration_test
label_desc: Use image node_id 3 for calling the TestExtractClusterWorkflow obj's underlying member method code with example usage: obj.run_cluster_assignment(with_fsdp) without return types