PyTorch PINNs that actually train fast

What it does PINNs-Torch is a PyTorch implementation of Physics-Informed Neural Networks. It uses PyTorch Lightning and Hydra to structure training, meaning you can launch pre-built PDE experiments—such as Navier-Stokes or Schrödinger—from config files rather than writing boilerplate loops.

The interesting bit The authors didn’t just port the code; they chased speed. By locking in CUDA Graphs and TorchScript JIT compilation, they report training up to 9× faster than the original TensorFlow v1 implementation. That optimization is the actual focus of their NeurIPS 2023 paper.

Key highlights

• Built on PyTorch Lightning and Hydra, so experiments are config-driven and modular.
• Includes working examples for Navier-Stokes and Schrödinger, with more in the examples folder.
• Claims up to 9× speed-up over the original TensorFlow v1 PINNs code via CUDA Graphs and JIT.
• Published at DLDE III, NeurIPS 2023, with parallel TensorFlow v2 and JAX versions available.
• BSD-3 licensed.

Caveats

• The authors note this is a first release, so expect rough edges and potential bugs.
• Data for examples is fetched from a remote server automatically, which means offline runs aren’t fully self-contained.

Verdict Worth a look if you’re doing PDE research in PyTorch and need a faster PINN baseline than the legacy TensorFlow v1 reference. Skip it if you need a mature, battle-tested framework or fully offline datasets.