Track grad norm with multiple losses

If I have a multi-task setting where my loss is a sum of various weighted objectives, is there a convenient way for me to track the gradient norms of each objective in Lightning individually? Perhaps some tweak on track_grad_norm ?

Hello, my apology for the late reply. We are slowly converging to deprecate this forum in favor of the GH build-in version… Could we kindly ask you to recreate your question there - Lightning Discussions