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

Signed-off-by: Kirthi Shankar Sivamani <[email protected]>
NVIDIA · Jan 19, 2024 · 96b31f8 · 96b31f8
1 parent 4f9662f
commit 96b31f8
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,