Well I found the solution to my hassle. In case anybode else runs into this problem, here is my solution. I replaced the
zero_tensor[:, module.keep_idxs.to(o_tensor.device), :, :] = o_tensor
line with the following:
zero_tensor.index_copy_(1, module.keep_idxs, o_tensor)
The reason for the failure is the following: The onnx exporter seems not to support colon notation. At least not in PyTorch 1.3.1. The new line is supported according to the ONNX export list and it also works now.
Best regards