Merge branch 'main' into automatic_ensemble_new

graphnet-team · May 28, 2024 · de7e10b · de7e10b
2 parents 0932394 + fc5d955
commit de7e10b
Show file tree

Hide file tree

Showing 9 changed files with 682 additions and 465 deletions.
diff --git a/requirements/torch_cu118.txt b/requirements/torch_cu118.txt
@@ -1,4 +1,5 @@
 # Contains packages requirements for GPU installation
 --find-links https://download.pytorch.org/whl/torch_stable.html
 torch==2.2.0+cu118
+torchvision==0.17.0+cu118
 --find-links https://data.pyg.org/whl/torch-2.2.0+cu118.html
diff --git a/requirements/torch_cu121.txt b/requirements/torch_cu121.txt
@@ -1,4 +1,5 @@
 # Contains packages requirements for GPU installation
 --find-links https://download.pytorch.org/whl/torch_stable.html
 torch==2.2.0+cu121
+torchvision==0.17.0+cu121
 --find-links https://data.pyg.org/whl/torch-2.2.0+cu121.html
diff --git a/src/graphnet/models/components/embedding.py b/src/graphnet/models/components/embedding.py
@@ -93,6 +93,7 @@ def __init__(
                 f"{n_features} features."
             )
         elif n_features >= 6:
+
             hidden_dim = 6 * seq_length
         else:
             hidden_dim = int((n_features + 0.5) * seq_length)