Hi,
I have the following problem. I am trying to propagate multiple outputs out of my network which are scalars, for example, latency or memory consumption of respective layers in addition to the output itself. These outputs I would then like to add to the main loss, let’s say cross-entropy.
With a single GPU, I am using a @dataclass to accumulate the respective scalar layer outputs in an accumulator that I then add to the loss, which contents I add to the main loss. However, I do have multiple GPUs that I could utilise for training and I am not sure how to propage the respective scalars and combine them such that I could use .backward(). Any help is much appreciated. Thanks.