Torch Version UNetEncoderBlock causes multi-card training error

[Yeeesir]

When I use maxim_pytorch provided by link, and try multi-gpu training, the following error occurred

  File "/home/miniconda3/lib/python3.8/site-packages/torch/autograd/__init__.py", line 154, in backward
    Variable._execution_engine.run_backward(
RuntimeError: Function BroadcastBackward returned an invalid gradient at index 76 - got [0] but expected shape compatible with [0, 32, 2, 2]

I found that the key problem lies in the torch implementation of UNetEncoderBlock, other network structures did not introduce errors