My output is prediction with size (bach_size, 15) and I want my loss like:
a * BCE(prediction[:, 14], taret[:, 14]) + BCE(prediction[:, 0:14], taret[:, 0:14]) * (prediction[:, 14] == taret[:, 14] )
To be more specific, I want to use the highest bit as an indicator, so the loss has two parts. First part is the loss of the indicator bit, thus BCE(prediction[:, 14], taret[:, 14]). The second part is loss of lower 14 bits when indicator bit is right, and it is BCE(prediction[:, 0:14], taret[:, 0:14]) * (prediction[:, 14] == taret[:, 14] ).
How to implement it in Pytorch?