diff --git a/tools/pao-ml/pao-retrain.py b/tools/pao-ml/pao-retrain.py index 1cd57b03c9..2c30ee6914 100755 --- a/tools/pao-ml/pao-retrain.py +++ b/tools/pao-ml/pao-retrain.py @@ -45,7 +45,7 @@ def main() -> None: train_model(model, dataset, epochs=args.epochs, batch_size=args.batch) # Save the model. - torch.jit.save(model, args.model) # type: ignore + torch.jit.save(model, args.model) print(f"Saved model to file: {args.model}") diff --git a/tools/pao-ml/pao-train.py b/tools/pao-ml/pao-train.py index 5b74760ddd..7e455ff8e0 100755 --- a/tools/pao-ml/pao-train.py +++ b/tools/pao-ml/pao-train.py @@ -79,7 +79,7 @@ def main() -> None: # Save the model. default_fn = f"{dataset.kind.prim_basis_name}-PAO{dataset.kind.pao_basis_size}-{args.kind}.pt" fn = args.model or default_fn - torch.jit.save(model, fn) # type: ignore + torch.jit.save(model, fn) print(f"Saved model to file: {fn}") diff --git a/tools/pao-ml/requirements.txt b/tools/pao-ml/requirements.txt index b0c38aa07b..dac214f897 100644 --- a/tools/pao-ml/requirements.txt +++ b/tools/pao-ml/requirements.txt @@ -1,6 +1,77 @@ +aiohappyeyeballs==2.6.1 +aiohttp==3.13.3 +aiosignal==1.4.0 +antlr4-python3-runtime==4.9.3 +ase==3.27.0 +attrs==25.4.0 +certifi==2026.1.4 +charset-normalizer==3.4.4 +contourpy==1.3.3 +cuda-bindings==12.9.4 +cuda-pathfinder==1.3.4 +cycler==0.12.1 +e3nn==0.5.9 +filelock==3.24.2 +fonttools==4.61.1 +frozenlist==1.8.0 +fsspec==2026.2.0 +hydra-core==1.3.2 +idna==3.11 +Jinja2==3.1.6 +kiwisolver==1.4.9 +librt==0.8.1 +lightning==2.6.1 +lightning-utilities==0.15.2 +lmdb==1.7.5 +MarkupSafe==3.0.3 +matplotlib==3.10.7 +matscipy==1.2.0 +mpmath==1.3.0 +multidict==6.7.1 +mypy==1.19.0 +mypy_extensions==1.1.0 nequip==0.16.1 -torch -numpy -e3nn +networkx==3.6.1 +numpy==2.3.5 +nvidia-cublas-cu12==12.8.4.1 +nvidia-cuda-cupti-cu12==12.8.90 +nvidia-cuda-nvrtc-cu12==12.8.93 +nvidia-cuda-runtime-cu12==12.8.90 +nvidia-cudnn-cu12==9.10.2.21 +nvidia-cufft-cu12==11.3.3.83 +nvidia-cufile-cu12==1.13.1.3 +nvidia-curand-cu12==10.3.9.90 +nvidia-cusolver-cu12==11.7.3.90 +nvidia-cusparse-cu12==12.5.8.93 +nvidia-cusparselt-cu12==0.7.1 +nvidia-nccl-cu12==2.27.5 +nvidia-nvjitlink-cu12==12.8.93 +nvidia-nvshmem-cu12==3.4.5 +nvidia-nvtx-cu12==12.8.90 +omegaconf==2.3.0 +opt-einsum-fx==0.1.4 +opt_einsum==3.4.0 +packaging==26.0 +pathspec==1.0.4 +pillow==12.1.1 +propcache==0.4.1 +pyparsing==3.3.2 +python-dateutil==2.9.0.post0 +pytorch-lightning==2.6.1 +PyYAML==6.0.3 +requests==2.32.4 +scipy==1.17.0 +setuptools==82.0.0 +six==1.17.0 +sympy==1.14.0 +torch==2.10.0 +torchmetrics==1.8.2 +tqdm==4.67.3 +triton==3.6.0 +types-requests==2.32.4.20250913 +typing_extensions==4.15.0 +urllib3==2.6.3 +yarl==1.22.0 + #EOF