Update number of layers and hidden dim (requires halving batch size t…

…o 16)
JustinDoIt · Sep 15, 2023 · 50a2239 · 50a2239
1 parent 54ce6c7
commit 50a2239
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/proteinworkshop/config/encoder/tfn.yaml b/proteinworkshop/config/encoder/tfn.yaml
@@ -2,8 +2,8 @@ _target_: proteinworkshop.models.graph_encoders.tfn.TensorProductModel
 r_max: 10.0 # Maximum distance for radial basis functions
 num_basis: 16 # Number of radial basis functions
 max_ell: 2 # Maximum degree/order of spherical harmonics basis functions and node feature tensors
-num_layers: 4 # Number of layers in the model
-hidden_irreps: "64x0e + 8x0o + 8x1e + 8x1o + 4x2e + 4x2o" # Irreps string for intermediate layer node feature tensors
+num_layers: 6 # Number of layers in the model
+hidden_irreps: "64x0e + 64x0o + 8x1e + 8x1o + 4x2e + 4x2o" # Irreps string for intermediate layer node feature tensors
 mlp_dim: 256 # Dimension of MLP for computing tensor product weights
 aggr: "sum" # Aggregation function to use
 pool: "mean" # Pooling operation to use