Add PyG packages to torch backend (#15911)

## Summary

These are now supported on pyx.
This commit is contained in:
Charlie Marsh 2025-09-17 10:18:30 -04:00 committed by GitHub
parent d4806ee921
commit 48f507680c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 2 deletions

View File

@ -322,11 +322,16 @@ impl TorchStrategy {
TorchSource::Pyx => { TorchSource::Pyx => {
matches!( matches!(
package_name.as_str(), package_name.as_str(),
"flash-attn" "deepspeed"
| "flash-attn"
| "flash-attn-3" | "flash-attn-3"
| "megablocks" | "megablocks"
| "natten" | "natten"
| "deepspeed" | "pyg-lib"
| "torch-cluster"
| "torch-scatter"
| "torch-sparse"
| "torch-spline-conv"
| "vllm" | "vllm"
| "torch" | "torch"
| "torch-model-archiver" | "torch-model-archiver"