PAO-ML: Pin versions of all Python packages

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}")
 
 

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}")
 
 

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