transformer_engine/pytorch/utils.py · 5b6ef0545ed14497ff296d18c3618ace4ba391db · OpenDAS / TransformerEngine

Enforce PyTorch version 2.1 and run attention tests with torch.compile (#1516) · 303c6d16

Kirthi Shankar Sivamani authored Feb 28, 2025



* Enforce torch 2.0 and run attn tests with torch.compile
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* replace torch.compile with jit_fuser
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

303c6d16

utils.py 12.1 KB

Replace utils.py