text: import argparse
  
  import pytorch_lightning as pl
  import torch
  from pytorch_lightning import Trainer
  from solo.methods import SimCLR
  
  from .utils import DATA_KWARGS, gen_base_kwargs, gen_batch, prepare_dummy_dataloaders
  
  
  def test_simclr():
      method_kwargs = {
          "output_dim": 256,
          "proj_hidden_dim": 2048,
          "temperature": 0.2,
          "supervised": False,
      }
  
      BASE_KWARGS = gen_base_kwargs(cifar=False)
      kwargs = {**BASE_KWARGS, **DATA_KWARGS, **method_kwargs}
      model = SimCLR(**kwargs)
  
      batch, batch_idx = gen_batch(
          BASE_KWARGS["batch_size"], BASE_KWARGS["num_classes"], "imagenet100"
      )
      loss = model.training_step(batch, batch_idx)
  
      assert loss != 0
  
      BASE_KWARGS = gen_base_kwargs(cifar=True)
      kwargs = {**BASE_KWARGS, **DATA_KWARGS, **method_kwargs}
      model = SimCLR(**kwargs)
  
      batch, batch_idx = gen_batch(BASE_KWARGS["batch_size"], BASE_KWARGS["num_classes"], "cifar10")
      loss = model.training_step(batch, batch_idx)
  
      assert loss != 0
  
      method_kwargs = {
          "output_dim": 256,
          "proj_hidden_dim": 2048,
          "temperature": 0.2,
          "supervised": True,
      }
  
      BASE_KWARGS = gen_base_kwargs(cifar=False)
      kwargs = {**BASE_KWARGS, **DATA_KWARGS, **method_kwargs}
      model = SimCLR(**kwargs)
  
      batch, batch_idx = gen_batch(
          BASE_KWARGS["batch_size"], BASE_KWARGS["num_classes"], "imagenet100"
      )
      loss = model.training_step(batch, batch_idx)
  
      assert loss != 0
  
      # test arguments
      parser = argparse.ArgumentParser()
      parser = pl.Trainer.add_argparse_args(parser)
      assert model.add_model_specific_args(parser) is not None
  
      # test parameters
      assert model.learnable_params is not None
  
      out = model(batch[1][0])
      assert (
          "logits" in out
          and isinstance(out["logits"], torch.Tensor)
          and out["logits"].size() == (BASE_KWARGS["batch_size"], BASE_KWARGS["num_classes"])
      )
      assert (
          "feats" in out
          and isinstance(out["feats"], torch.Tensor)
          and out["feats"].size() == (BASE_KWARGS["batch_size"], model.features_dim)
      )
      assert (
          "z" in out
          and isinstance(out["z"], torch.Tensor)
          and out["z"].size() == (BASE_KWARGS["batch_size"], method_kwargs["output_dim"])
      )
  
      args = argparse.Namespace(**kwargs)
      trainer = Trainer.from_argparse_args(
          args,
          checkpoint_callback=False,
          limit_train_batches=2,
          limit_val_batches=2,
      )
      train_dl, val_dl = prepare_dummy_dataloaders(
          "imagenet100", BASE_KWARGS["num_crops"], BASE_KWARGS["num_classes"], multicrop=False
      )
      trainer.fit(model, train_dl, val_dl)
  
      # normal training
      BASE_KWARGS = gen_base_kwargs(cifar=False, multicrop=False)
      kwargs = {**BASE_KWARGS, **DATA_KWARGS, **method_kwargs}
      model = SimCLR(**kwargs)
  
      args = argparse.Namespace(**kwargs)
      trainer = Trainer.from_argparse_args(
          args,
          checkpoint_callback=False,
          limit_train_batches=2,
          limit_val_batches=2,
      )
      train_dl, val_dl = prepare_dummy_dataloaders(
          "imagenet100",
          num_crops=BASE_KWARGS["num_crops"],
          num_small_crops=0,
          num_classes=BASE_KWARGS["num_classes"],
          multicrop=False,
      )
      trainer.fit(model, train_dl, val_dl)
  
      # multicrop
      BASE_KWARGS = gen_base_kwargs(cifar=False, multicrop=True, num_small_crops=6)
      kwargs = {**BASE_KWARGS, **DATA_KWARGS, **method_kwargs}
      model = SimCLR(**kwargs)
  
      args = argparse.Namespace(**kwargs)
      trainer = Trainer.from_argparse_args(
          args,
          checkpoint_callback=False,
          limit_train_batches=2,
          limit_val_batches=2,
      )
      train_dl, val_dl = prepare_dummy_dataloaders(
          "imagenet100",
          num_crops=BASE_KWARGS["num_crops"],
          num_small_crops=6,
          num_classes=BASE_KWARGS["num_classes"],
          multicrop=True,
      )
      trainer.fit(model, train_dl, val_dl)
  