cleanup unit offset again

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'x-transformers',
   packages = find_packages(exclude=['examples']),
-  version = '1.31.6',
+  version = '1.31.8',
   license='MIT',
   description = 'X-Transformers - Pytorch',
   author = 'Phil Wang',

x_transformers/x_transformers.py

@@ -580,7 +580,7 @@ def __init__(
 
     def forward(self, x):
         normed = self.ln(x)
-        gamma = self.gamma + self.unit_offset
+        gamma = self.gamma + float(self.unit_offset)
         return normed * gamma
 
 class AdaptiveLayerNorm(Module):
@@ -615,7 +615,8 @@ def __init__(
         nn.init.constant_(self.g, 1. - float(unit_offset))
 
     def forward(self, x):
-        return F.normalize(x, dim = -1) * self.scale * (self.g + self.unit_offset)
+        gamma = self.g + float(self.unit_offset)
+        return F.normalize(x, dim = -1) * self.scale * gamma
 
 class RMSNorm(Module):
     def __init__(
@@ -631,7 +632,8 @@ def __init__(
         nn.init.constant_(self.g, 1. - float(unit_offset))
 
     def forward(self, x):
-        return F.normalize(x, dim = -1) * self.scale * (self.g + self.unit_offset)
+        gamma = self.g + float(self.unit_offset)
+        return F.normalize(x, dim = -1) * self.scale * gamma
 
 class AdaptiveRMSNorm(Module):
     def __init__(
@@ -747,11 +749,13 @@ def __init__(
     def forward(self, x, **kwargs):
         out = self.fn(x, **kwargs)
 
+        gamma = self.gamma + float(self.unit_offset)
+
         if isinstance(out, Tensor):
-            return out * self.gamma
+            return out * gamma
 
         out, *rest = out
-        return out * self.gamma, *rest
+        return out * gamma, *rest
 
 class AdaptiveLayerScale(Module):
     def __init__(