I had to implement my own loss function which specified custom gradients as well by normalizing some regularization term’s gradients (this case the non-informative prior for geodesic distance + L1 loss of the in-plane rotation) accordingly to the softmax conditioned outputs.
You can specify such custom forward and backwards procedure using PyTorch Functions or Modules.
http://pytorch.org/docs/master/notes/extending.html 12 demonstrates how you can attain a gradient input and manipulate it in this case to replicate the backwards procedure shown in the Caffe code.
If your loss simply requires functional differentiation, then you can just create a nn.Module and have the auto-diff handle it for you :).
An example of it is available in some of my bundle of code here for a structured Mahalanobis metric loss.

  
      github.com
  
  
    iwasaki-kenta/keita/blob/master/losses/mahalanobis.py 147
"""
Minimizing Mahalanobis distance between related pairs, and maximizing between negative pairs.

A loss typically used for creating a Euclidian embedding space for a wide variety of supervised learning problems.
The original implementation was by Davis King @ Dlib.

PyTorch Implementation: https://gist.github.com/bkj/565c5e145786cfd362cffdbd8c089cf4

Made changes such that accuracy is provided on a forward pass as well.
"""

import torch
import torch.nn.functional as F
from torch import nn


class MahalanobisMetricLoss(nn.Module):
    def __init__(self, margin=0.6, extra_margin=0.04):
        super(MahalanobisMetricLoss, self).__init__()



  This file has been truncated. show original