Thanks for the repro instructions.
This error is raised when executing it:
Epoch 1
  0%|                                                                                                                                                                          | 0/2500 [00:01<?, ?batch/s]
Traceback (most recent call last):
  File "ssl_models/src/main.py", line 96, in <module>
    main(args)
  File "ssl_models/src/main.py", line 53, in main
    trainer.fit()
  File "/workspace/src/ssl_models/src/trainer.py", line 51, in fit
    train_epoch_loss = self.train_epoch(self.model, self.train_loader, self.optimizer)
  File "/workspace/src/ssl_models/src/trainer.py", line 83, in train_epoch
    train_loss = model.training_step(batch, step)
  File "/workspace/src/ssl_models/src/simclr.py", line 115, in training_step
    loss = self.criterion(x1,x2)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1056, in _call_impl
    return forward_call(*input, **kwargs)
  File "/workspace/src/ssl_models/src/losses.py", line 32, in forward
    sim = self.similarity_f(z.unsqueeze(1), z.unsqueeze(0)) / self.temperature
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1056, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/distance.py", line 75, in forward
    return F.cosine_similarity(x1, x2, self.dim, self.eps)
RuntimeError: cosine_similarity requires both inputs to have the same sizes, but x1 has [32, 1, 128] and x2 has [1, 32, 128]