Merge pull request #501 from datamol-io/Patch-for-mup

Update global_architectures.py
datamol-io · Feb 16, 2024 · 48d26fd · 48d26fd
2 parents 7771164 + 0e6c932
commit 48d26fd
Showing 1 changed file with 0 additions and 4 deletions.
diff --git a/graphium/nn/architectures/global_architectures.py b/graphium/nn/architectures/global_architectures.py
@@ -1339,11 +1339,7 @@ def _recursive_divide_dim(x: collections.abc.Mapping):
                 elif k in ["in_dim", "out_dim", "in_dim_edges", "out_dim_edges"]:
                     x[k] = round(v / divide_factor)
                 elif k in ["embed_dim"]:
-                    num_heads = x.get("num_heads", 1)
                     x[k] = round(v / divide_factor)
-                    assert (
-                        x[k] % num_heads == 0
-                    ), f"embed_dim={x[k]} is not divisible by num_heads={num_heads}"
 
         _recursive_divide_dim(kwargs["layer_kwargs"])