Bug with 3D convolutions

[EiffL]

We got some weird deadlock when trying to run a simple 3d conv with blocks, most likely from the halo exchange:

all_inspector.cc:107] One or more tensors were submitted to be reduced, gathered or broadcasted by subset of ranks and are waiting for remainder of ranks for more than 60 seconds. This may indicate that different ranks are trying to submit different tensors or that only subset of ranks is submitting tensors, which will cause deadlock. 
Missing ranks:
1: [shift/HorovodAllgather_shift_ExpandDims_0]
3: [shift_1/HorovodAllgather_shift_1_ExpandDims_0]

@b-remy can you document here exactly how this happened? We'll need to sort it out....

[tobias-liaudat]

Leaving two comments as I'm online right now.

To reproduce the error one can run this script with this job.

After finding the deadlock, we changed the 3D convolution by a dense layer and it was working fine in this test script and this job.