That error looks like it’s coming from the actual operation, not anything to do with tracing. Are you tracing it with the same inputs you’re using to run it successfully in normal PyTorch?
If your model has any control flow that depends on the inputs tracing may not work correctly. Can you also post the full code you’re using to run/trace your model?