If I use torch.jit.script instead of torch.jit.trace:
sm = torch.jit.script(wrapper)
torch.jit.save(sm, "traced_model.pt")

I realized I could put print statements and that they would show up in the c++ console. I put some surrounding the execution of torch.stack in my ModelWrapper. It turns out this is the moment it throws the exception. Same thing happens when trying to write it with torch.cat followed by torch.unsqueeze (the cat fails).
Weirdly I am able to do this inside the ModuleWrapper’s forward:
something = torch.stack([torch.randn([2, 3, 4]), torch.randn([2, 3, 4])])
When I removed the torch.stack call from ModelWrapper, using print statements I was able to pinpoint a failure within the u2netmodel forward on a call to cat. https://github.com/NathanUA/U-2-Net/blob/b77cd6da3204efcb03e18e15dd3b9eb24d47f969/model/u2net.py#L87
So my question now is why can’t I do torch.stack or torch.cat in this module? Is it because these things allocate memory whereas calls to Conv2D don’t allocate? I’ve used print statements to make sure everything’s on cuda etc. Something about cloning? Python execution of the same model is working fine.