I am trying to implement the Geometric-Aware loss function outline in RenderForCNN (https://arxiv.org/pdf/1505.05641.pdf 49) whose caffe implementation can be found here (https://github.com/charlesq34/caffe-render-for-cnn/blob/view_prediction/src/caffe/layers/softmax_view_loss_layer.cpp 19).
The problem is that since the output is being done to multiple labels at the same time, one has to only backprop to a specific number of parameters, this is done by assigning specific weights to each gradient, with many of them being 0 (lines 139 and 177-187 in file linked above). Is there a way to implement a similar fuction in PyTorch ? if so, any advice on how to do this ? Thank you!