A linter for PyTorch 'torch-preflight' [P]

Been working on this for the last few months. I've been working on PyTorch for the past few years and I always felt, many a times my work went into dump, because of some mistakes I made in the code. torch-preflight reads your PyTorch code and catches the bugs costing you GPU hours. Things like losses.append(loss), which holds the autograd graph from every step until CUDA dies on you or no zero_grad() in the loop or gradient accumulation without dividing the loss or DDP with no DistributedSampler, so every rank trains on the same batches. I've been able to get 13 rules so far. Your code never gets imported or executed, so you need no GPU and no torch install. There's another part to this that estimates VRAM. Point the tool at a training script and a GPU, and you learn whether the run fits before you pay for the instance. You also get the list of changes to make the run fit, with the GiB each one saves. pip install torch-preflight https://github.com/highwaterlabs/torch-preflight https://pypi.org/project/torch-preflight/ Please try this out, and I would like to get your feedback! It's still a work in progeress. Would like to know what breaks on your code. False positives kill a linter, and my only large test target so far has been the PyTorch source tree. Same for the memory numbers. Mine land within 4% of measured peaks, but from four models on one T4. PS: open to contributions, and issues are already open on the repo. Soon I'm going to add a few "Good first issues" as well. Feel free to ping me if you have any questions! submitted by /u/LeJanbandhu [link] [comments]

Source: r/MachineLearning | 2026-08-14