BYOL loss function

[vishnu-dev]

🐛 Bug

The loss function in BYOL doesn't seem to match the one defined in the paper.

To Reproduce

Running the training for BYOL.

Code sample

lightning-bolts/src/pl_bolts/models/self_supervised/byol/byol_module.py

Lines 130 to 142 in 4f910f6

	def calculate_loss(self, v_online: Tensor, v_target: Tensor) -> Tensor:
	"""Calculates similarity loss between the online network prediction of target network projection.
	Args:
	v_online (Tensor): Online network view
	v_target (Tensor): Target network view
	"""
	_, z1 = self.online_network(v_online)
	h1 = self.predictor(z1)
	with torch.no_grad():
	_, z2 = self.target_network(v_target)
	return -2 * F.cosine_similarity(h1, z2).mean()

Expected behavior

The loss should be 2 - (2 * F.cosine_similarity(h1, z2).mean())?

[Borda]

Thank you, @vishnu-dev mind sending PR with the fix? 🐰