Fix trainer dataloader slow loading issue (#3219)

* Fix trainer dataloader handling in src/axolotl/core/trainers/base.py

* update comment to reflect torch version

---------

Co-authored-by: Wing Lian <wing.lian@gmail.com>
This commit is contained in:
Qingyang Wu
2025-10-22 07:22:14 -07:00
committed by GitHub
parent 613bcf90e5
commit 3750fdcf79
2 changed files with 13 additions and 12 deletions

View File

@@ -49,7 +49,7 @@ def parse_requirements(extras_require_map):
try:
torch_version = version("torch")
except PackageNotFoundError:
torch_version = "2.6.0" # default to torch 2.6
torch_version = "2.8.0" # default to torch 2.8.0
_install_requires.append(f"torch=={torch_version}")
version_match = re.match(r"^(\d+)\.(\d+)(?:\.(\d+))?", torch_version)