Rename and remove unused argument

Author: jacobbieker

graph_weather/models/aurora/model.py

@@ -188,12 +188,11 @@ def __init__(
         self,
         input_features: int,
         output_features: int,
-        embed_dim: int = 256,
         latent_dim: int = 256,
         num_layers: int = 4,
         max_points: int = 10000,
         max_seq_len: int = 1024,
-        use_checkPointing: bool = False,
+        use_checkpointing: bool = False,
     ):
         super().__init__()
 
@@ -203,12 +202,12 @@ def __init__(
         self.output_features = output_features
 
         # Model components
-        self.encoder = PointEncoder(input_features, embed_dim, max_seq_len)
-        self.processor = PointCloudProcessor(embed_dim, num_layers)
-        self.decoder = PointDecoder(embed_dim, output_features)
+        self.encoder = PointEncoder(input_features, latent_dim, max_seq_len)
+        self.processor = PointCloudProcessor(latent_dim, num_layers)
+        self.decoder = PointDecoder(latent_dim, output_features)
 
         # Add gradient checkpointing
-        self.use_checkpointing = use_checkPointing
+        self.use_checkpointing = use_checkpointing
 
         # Initialize weights properly
         self._init_weights()

tests/test_aurora.py

@@ -44,8 +44,7 @@ def model_config():
     return {
         "input_features": 2,
         "output_features": 2,
-        "embed_dim": 8,  # Adjusted to avoid overfitting
-        "latent_dim": 16,  # Adjusted
+        "latent_dim": 8,  # Adjusted
         "max_points": 50,  # Adjusted
         "max_seq_len": 128,  # Adjusted
     }
@@ -74,8 +73,7 @@ def test_gradient_checkpointing_config():
     config_no_checkpoint = {
         "input_features": 2,
         "output_features": 2,
-        "embed_dim": 8,
-        "latent_dim": 16,
+        "latent_dim": 8,
         "max_points": 50,
         "max_seq_len": 128,
         "use_checkpointing": False,