torch._assert() is traceable, you should be able to use instead of builtin assert. Additionally, for modules where you may not want/be able to modify the source to use torch._assert(), there is an AST Rewriter tracer in FX experimental, see here 3.