code: def test_swav_multi_prototypes(self, device: str = "cpu", seed=0):
      for in_features, _, n_prototypes in self.swavProtoypes:
          torch.manual_seed(seed)
          prototypes = SwaVPrototypes(in_features, n_prototypes)
          prototypes = prototypes.eval()
          prototypes = prototypes.to(device)
          for batch_size in [1, 2]:
              msg = (
                  "prototypes d_in, n_prototypes = "
                  + f"{in_features} x {n_prototypes}"
              )
              with self.subTest(msg=msg):
                  x = torch.torch.rand((batch_size, in_features)).to(
                      device
                  )
                  with torch.no_grad():
                      y = prototypes(x)
                  for layerNum, prototypeSize in enumerate(
                      n_prototypes
                  ):
                      self.assertEqual(y[layerNum].shape[0], batch_size)
                      self.assertEqual(
                          y[layerNum].shape[1], prototypeSize
                      )
  
file_path: lightly/tests/models/test_ProjectionHeads.py
imports: import unittest
  import torch
  import lightly
  from lightly.models.modules.heads import BarlowTwinsProjectionHead, BYOLPredictionHead, BYOLProjectionHead, DINOProjectionHead, MoCoProjectionHead, MSNProjectionHead, NNCLRPredictionHead, NNCLRProjectionHead, SimCLRProjectionHead, SimSiamPredictionHead, SimSiamProjectionHead, SwaVProjectionHead, SwaVPrototypes, TiCoProjectionHead, VicRegLLocalProjectionHead
label_desc: Use image node_id 6 for calling the TestProjectionHeads obj's underlying member method code with example usage: obj.test_swav_multi_prototypes(device, seed) without return types