For such problem, it will definitely be much easier to let the autograd figure out the gradients. In particular because you might not have all the information you need during the backward of that layer as the loss will depend on the layers that appeared after this one.
Note as well that a good ressource is other implementations of contrastive losses for pytorch: https://github.com/topics/contrastive-loss 2