I think you might try your luck with torch.fx.wrap described in the non-torch functions section of the fx documentation 3.
Best regards
Thomas
P.S.: So the more detailed explanation for the error is that during tracing, the function will be called not with a tensor x but with a proxy object that causes functions to amend to the trace. wrap is a mechanism make function calls with these proxy objects insert a record of their invocation.