Hi,
I am not sure to understand exactly what you want to do here. But if you want to add special penalties to your loss, you can simply add them before calling backward:
loss = crit(out, target)
final_loss = loss + penalty(out, target)
final_loss.backward()