overrideable dropout rate for fine-tuning

Author: lucidrains

README.md

@@ -138,7 +138,7 @@ You can also load, with overriding of the `target_length` parameter, if you are
 ```python
 from enformer_pytorch import load_pretrained_model
 
-model = load_pretrained_model('preview', target_length = 128)
+model = load_pretrained_model('preview', target_length = 128, dropout_rate = 0.1)
 
 # do your fine-tuning
 ```
@@ -148,7 +148,7 @@ You can also define the model externally, and then load the pretrained weights b
 ```python
 from enformer_pytorch import Enformer, load_pretrained_model
 
-enformer = Enformer(dim = 1536, depth = 11, target_length = 128)
+enformer = Enformer(dim = 1536, depth = 11, target_length = 128, dropout_rate = 0.1)
 
 load_pretrained_model('preview', model = enformer)
 

enformer_pytorch/model_loader.py

@@ -36,6 +36,7 @@ def load_pretrained_model(
     slug,
     force = False,
     target_length = None,
+    dropout_rate = None,
     model = None
 ):
     if slug not in CONFIG:
@@ -57,7 +58,7 @@ def load_pretrained_model(
 
     # load
 
-    override_params = remove_nones({'target_length': target_length})
+    override_params = remove_nones({'target_length': target_length, 'dropout_rate': dropout_rate})
     params = {**config['params'], **override_params}
 
     if not exists(model):

setup.py

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