create nn.py

basedformer/gptj.py

 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from basedformer.utils import *
 from torch.utils.checkpoint import checkpoint as ck
 from einops import rearrange, repeat
 try:

basedformer/nn.py

+# extension to torch.nn or torch.nn.functional.
+# for neural network/like ops.
+import torch
+import math
+
+def gelu_new(x):
+    return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
\ No newline at end of file