From 72c71b67ec268b29d1d33c83be421c78f9c5bcb0 Mon Sep 17 00:00:00 2001
From: luoling-LC <105470086+luoling-LC@users.noreply.github.com>
Date: Fri, 13 May 2022 18:03:11 +0800
Subject: [PATCH] [NFC] polish colossalai/kernel/jit/bias_gelu.py code style
 (#946)

Co-authored-by: jnbai <897086360@qq.com>
---
 colossalai/kernel/jit/bias_gelu.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/colossalai/kernel/jit/bias_gelu.py b/colossalai/kernel/jit/bias_gelu.py
index f7a425dd5..e6da70c40 100644
--- a/colossalai/kernel/jit/bias_gelu.py
+++ b/colossalai/kernel/jit/bias_gelu.py
@@ -1,6 +1,5 @@
 import torch
 
-
 ###### BIAS GELU FUSION/ NO AUTOGRAD ################
 # 1/sqrt(2*pi)-> 0.3989423
 # 1/sqrt(2)   -> 0.70710678
@@ -9,10 +8,12 @@ import torch
 # actual gelu is:
 # x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
 
+
 @torch.jit.script
 def bias_gelu(bias, y):
     x = bias + y
-    return  x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
+    return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
+
 
 # gradient of tanh approximation of gelu
 # gradient of actual gelu is:
@@ -23,9 +24,11 @@ def bias_gelu_back(g, bias, y):
     tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
     # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
     ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)
-    return ff*g
+    return ff * g
+
 
 class GeLUFunction(torch.autograd.Function):
+
     @staticmethod
     # bias is an optional argument
     def forward(ctx, input, bias):
@@ -38,4 +41,5 @@ class GeLUFunction(torch.autograd.Function):
         tmp = bias_gelu_back(grad_output, bias, input)
         return tmp, tmp
 
-bias_gelu_impl = GeLUFunction.apply
\ No newline at end of file
+
+bias_gelu_impl = GeLUFunction.apply