Maybe a small bug in LARS implementation

[Big-Brother-Pikachu]

Hello, thanks for your pretty implementation. I think I may find a small bug in your LARS implementation.trust_ratio = tf.where( tf.greater(w_norm, 0), tf.where( tf.greater(g_norm, 0), (self.eeta * w_norm / g_norm), 1.0), 1.0) is a little different from https://github.com/Spijkervet/SimCLR/blob/654f05f107ce17c0a9db385f298a2dc6f8b3b870/modules/lars.py#L119-L127 As greater is > and ge is >=. Thus bias paramater which is initialized as 0 is never updated. I think trust_ratio = torch.where( w_norm.gt(0), torch.where( g_norm.gt(0), (self.eeta * w_norm / g_norm), torch.Tensor([1.0]).to(device), ), torch.Tensor([1.0]).to(device), ).item() may work better.