Saurav_Gupta1:

Am I gathering loss correctly?


It looks correct to me.

I am asking because in the above example tmp contains loss from GPU:0 but in some other run tmp contains loss of GPU:1

Could you please share a repro where tmp[0] contains loss from rank 1? The above output seems only shows loss from rank 0?