ONNX export failed with broadcasted operation

Hi, I’m new to PyTorch, helping out with a friend’s project.

Trying to export a model to onnx, but encountering this

RuntimeError: ONNX export failed: Could not export a broadcasted operation; ONNX likely does not support this form of broadcasting.

I have troubleshooted the problematic line to a line that is doing
z = function(z) * z + x
where z.size() and x.size() are both torch.Size([1, 16, 40, 63])
and function(z).size() is torch.Size([1, 16, 1, 1])

So I tried doing
z = function(z).expand_as(z) * z + x
or
z = function(z).expand(z.size()) * z + x

But both alternatives also give the same error Could not export a broadcasted operation; ONNX likely does not support this form of broadcasting.

Any suggestions to get around this issue?