Avoid using torch.compile for roll and fill_ (#609)

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
NVIDIA · Jan 19, 2024 · 5e80b2a · 5e80b2a
1 parent b4b8ae7
commit 5e80b2a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/transformer_engine/pytorch/fp8.py b/transformer_engine/pytorch/fp8.py
@@ -583,7 +583,7 @@ def _update_amax_history(amax_history: torch.Tensor) -> torch.Tensor:
     return amax_history
 
 
-@jit_fuser
+@torch.jit.script
 def _default_get_amax(
     amax_history: torch.Tensor,
     amax_compute_algo: str,
@@ -625,7 +625,7 @@ def _compute_scaling_factor_inverse(
     return torch.where(non_weight_mask, 1.0 / scale, scale_inv)
 
 
-@jit_fuser
+@torch.jit.script
 def _fused_amax_and_scale_update(
     amax_history: torch.Tensor,
     scale: torch.Tensor,