I've been looking at using compilation in torch but not successful yet as otherwise it can take awhile to run. https://pytorch.org/tutorials/intermediate/torch_compile_tut...