[NFC] polish colossalai/kernel/jit/bias_gelu.py code style (#946)

Co-authored-by: jnbai <897086360@qq.com>

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