In doing so, this will happen as mentioned by @ptrblck. So, the layers before the custom layers will also have gradients of contrastive loss accumulated in addition to that of CE loss. But, I want to restrict this gradient accumulation (of contrastive + CE loss) within the custom layer.
Probably, this would justify why I am doing it this way. ![]()