Solved by tjoseph in post #3 
              
              
                For everyone wondering: requires_grad is not supposed to work. trace only tracks tensor operations, not attributes. See [JIT] jit.trace does not support parameter.requires_grad? · Issue #53515 · pytorch/pytorch · GitHub 
Anyone know a solution? This is something that is supported in tensorflow, but …