CI success (#64)

* CI success

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

metnet/models/metnet.py

@@ -41,7 +41,7 @@ def __init__(
         num_att_layers = self.config["num_att_layers"]
         output_channels = self.config["output_channels"]
         use_preprocessor = self.config["use_preprocessor"]
-        num_att_heads = self.config["num_att_heads"]
+        num_att_heads = self.config.get("num_att_heads", 16)
 
         self.forecast_steps = forecast_steps
         self.input_channels = input_channels

metnet/models/metnet_pv.py

@@ -53,7 +53,7 @@ def __init__(
         pv_fc_out_channels = self.config["pv_fc_out_channels"]
         pv_id_embedding_channels = self.config["pv_id_embedding_channels"]
         fc_1_channels = self.config["fc_1_channels"]
-        num_att_heads = self.config["num_att_heads"]
+        num_att_heads = self.config.get("num_att_heads", 16)
         avg_pool_size = self.config["avg_pool_size"]
 
         self.forecast_steps = forecast_steps
@@ -117,7 +117,20 @@ def __init__(
         )
 
         # FC layer, takes in 'econcded_timesteps', pv history, and embedding of pv ids
-        self.fc1 = nn.Linear(in_features=606976, out_features=fc_1_channels)
+        # self.fc1 = nn.Linear(in_features=606976, out_features=fc_1_channels)
+
+        # Calculating the number of features after average pooling and reshape
+        num_features_after_pooling = (input_size // 4 // avg_pool_size) ** 2 * hidden_dim
+
+        # Calculating the total number of features for the linear layer input
+        total_features = (
+            num_features_after_pooling
+            + (n_pv_timestamps * pv_fc_out_channels)
+            + (pv_id_embedding_channels * num_pv_systems)
+        )
+
+        # Updating the in_features parameter of nn.Linear
+        self.fc1 = nn.Linear(in_features=total_features, out_features=fc_1_channels)
 
     def encode_timestep(self, x, pv_yield_history, fstep=1):
         # Preprocess Tensor