Thank you so much for taking the time to respond!
You’re absolutely right — most of the logic can indeed live in higher-level libraries like PyTorch Lightning or Accelerate. However, the internal structure of this tool is highly modular and granular.
For example, even without using the decorators, core functions such as safe_to_device(), try_batch_size(), or run_with_amp() can be used individually — making them potential candidates for inclusion in torch.utils or torch.cuda.
To be clear: this is not a proposal to merge the whole package as-is. Rather, I wanted to offer a small set of reusable utilities that address common patterns like:
- clean device transfers
- OOM-safe function calls
- automatic AMP context
- fallback handling
- safe tensorization of inputs
These tools came out of repeated pain points in long training runs on low-memory devices, and I believe parts of them could help simplify day-to-day PyTorch code, even in core-level examples or docs.
If there’s any interest from the core team, I’d be happy to isolate and PR one or two focused components (e.g., safe_to_device() or a simple cuda_guard() context manager).
Thanks again!