directly set use_tf_gamma on Attention instances

Author: lucidrains

enformer_pytorch/modeling_enformer.py

@@ -480,6 +480,10 @@ def from_pretrained(name, use_tf_gamma = None, **kwargs):
     enformer = Enformer.from_pretrained(name, **kwargs)
 
     if name == 'EleutherAI/enformer-official-rough':
-        enformer.use_tf_gamma = default(use_tf_gamma, True)
+        use_tf_gamma = default(use_tf_gamma, True)
+
+        for module in enformer.modules():
+            if isinstance(module, Attention):
+                module.use_tf_gamma = use_tf_gamma
 
     return enformer

setup.py

@@ -4,7 +4,7 @@
   name = 'enformer-pytorch',
   packages = find_packages(exclude=[]),
   include_package_data = True,
-  version = '0.8.2',
+  version = '0.8.3',
   license='MIT',
   description = 'Enformer - Pytorch',
   author = 'Phil Wang',