Skip to content

Commit

Permalink
Apply isort and black reformatting
Browse files Browse the repository at this point in the history
Signed-off-by: ashvinnihalani <[email protected]>
  • Loading branch information
ashvinnihalani committed Jul 10, 2024
1 parent cd68086 commit e6940ed
Showing 1 changed file with 3 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -47,13 +47,13 @@ def bias_dropout_add_fused_train_(
# type: (Tensor, Tensor, Tensor, float) -> Tensor
return bias_dropout_add(x, bias, residual, prob, True)


@torch.jit.script
def dropout_add_fused_train_(
x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float
) -> torch.Tensor:
def dropout_add_fused_train_(x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float) -> torch.Tensor:
# type: (Tensor, None, Tensor, float) -> Tensor
return dropout_add(x, bias, residual, prob, True)


def bias_dropout_add_fused_train(x, bias, residual, prob):
# re-enable torch grad to enable fused optimization.
with torch.enable_grad():
Expand Down

0 comments on commit e6940ed

Please sign in to comment.