torch.compile() transforms a PyTorch function or Module instance into a compiled version intended to execute faster while behaving like the original during forward and backward passes. section 3.4.2
Compile a linear classifier
After defining a PyTorch linear classifier, compile the Module with compiled_model <- torch$compile(model). Compilation can also wrap a function directly, such as a dense operation using matrix multiplication and ReLU. section 3.4.2
input_dim <- 2L output_dim <- 1L W <- torch$rand(input_dim, output_dim, requires_grad = TRUE) b <- torch$zeros(output_dim, requires_grad = TRUE)
Use the associated code listing to locate the complete compiled-model call in the linear-classifier workflow. section 3.4.2
Use compilation with practical expectations:
- Treat it as optional: PyTorch commonly runs eagerly, and compilation may fail for some models or provide no speedup. section 3.4.2
- Recognize Dynamo as the compiler behind the API and account for its relatively recent, sometimes unreliable implementation. section 3.4.2 section 3.4.3
- Do not assume compilation closes PyTorch’s performance gap with JAX; large models may still run substantially faster with JAX. section 3.4.3
