Update dnaflash.py

dnaflash.py

@@ -118,12 +118,12 @@ class T5RelativePositionBias(nn.Module):
 class OffsetScale(nn.Module):
     def __init__(self, dim, heads = 1):
         super().__init__()
-        self.weight = nn.Parameter(torch.ones(heads, dim))
-        self.bias = nn.Parameter(torch.zeros(heads, dim))
-        nn.init.normal_(self.weight, std = 0.02)
+        self.gamma = nn.Parameter(torch.ones(heads, dim))
+        self.beta = nn.Parameter(torch.zeros(heads, dim))
+        nn.init.normal_(self.gamma, std = 0.02)
 
     def forward(self, x):
-        out = einsum('... d, h d -> ... h d', x, self.weight) + self.bias
+        out = einsum('... d, h d -> ... h d', x, self.gamma) + self.beta
         return out.unbind(dim = -2)
 
 # activation functions