Hi @agolynski,
Thanks for the reply. The goal is not to sync value ‘f’ across gpus and is conditioned upon the input mini-batch.
More context on what we are trying to do: We want to train a network which takes either a single tensor (‘A’ of shape Nx3x128x128) or pair (tensors ‘A’ and ‘B’ of shapes Nx3x128x128 and Mx3x32x32) as input and passed to NetA and NetB respectively. For example, tensor ‘A’ would be images and tensor ‘B’ could correspond to bounding box crops generated from an off the shelf object detector (in our case we do not alter or tune the parameters of the object detector). And if there are no objects detected in tensor ‘A’ then ‘B’ would be empty and we tried setting it to shape 0 zeros along the dimension 0. NetA and NetB do not share any parameters and are independent of each other.
One possibility we are aware of, is to make sure to run NetB only when tensor B can be obtained on all gpus. So, if we are running the model using DDP on 8 gpus, we could do the backward when tensor B is not empty on all 8 gpus, which is less frequent and would cause unreasonable amount of training time. We wanted to find a way to sync gradients only on gpus on which tensor B is not empty. Hope I am clear on what we are trying to do. Thanks.