directly set use_tf_gamma on Attention instances

lucidrains · Sep 28, 2023 · d2dbc21 · d2dbc21
1 parent ece4306
commit d2dbc21
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/enformer_pytorch/modeling_enformer.py b/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
diff --git a/setup.py b/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',