diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index bd3bc94b60..80386db2d9 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -213,6 +213,7 @@ def backward( # Restore saved tensors saved_tensors = restore_from_saved(func_ctx.tensor_objects, func_ctx.saved_tensors) + func_ctx.tensor_objects = None # Unflatten list of saved tensors for ctx in basic_op_ctxs: