Hi Howard,
thanks for helping me out!
Replacing nn.BatchNorm3D(out_channels) with nn.SyncBatchNorm(out_channels) worked.
I noticed that there is no nn.SyncBatchNorm3D? Do I need to take any extra measures here?
Best,
Thorsten