Fix gpu

climanet/dataset.py

@@ -238,7 +238,9 @@ def __getitem__(self, idx):
         monthly_t_patch = self.monthly_t[:, i : i + ph, j : j + pw]
 
         # (M, T, H, W) -> (M, T, pH, pW)
-        daily_nan_mask_t_patch = self.daily_nan_mask_t[:, :, i : i + ph, j : j + pw].unsqueeze(0)
+        daily_nan_mask_t_patch = self.daily_nan_mask_t[
+            :, :, i : i + ph, j : j + pw
+        ].unsqueeze(0)
 
         if self.land_mask_t is not None:
             land_t_patch = self.land_mask_t[i : i + ph, j : j + pw]  # (H, W)

climanet/predict.py

@@ -155,7 +155,7 @@ def predict_monthly_var(
     writer.add_scalar("Loss/Average", average_loss)
 
     if return_numpy:
-        all_predictions = all_predictions.numpy()
+        all_predictions = all_predictions.cpu().numpy()
 
     if save_predictions:
         if not return_numpy:

climanet/st_encoder_decoder.py

@@ -637,6 +637,7 @@ def __init__(
         dropout=0.0,
         sh_dim=96,
         scale_dim=10,
+        use_checkpoint=True,
     ):
         """Initialize the Spatio-Temporal Model.
 
@@ -663,6 +664,8 @@ def __init__(
             k: v for k, v in locals().items() if k not in ("self", "__class__")
         }
 
+        self.use_checkpoint = use_checkpoint
+
         self.encoder = VideoEncoder(
             in_chans=in_chans,
             embed_dim=embed_dim,
@@ -742,30 +745,50 @@ def forward(
         input_data_reshaped = input_data.reshape(B * M, C, T, H, W)
         daily_mask_reshaped = daily_mask.reshape(B * M, C, T, H, W)
 
-        latent = checkpoint(
-            self.encoder,
-            input_data_reshaped,
-            daily_mask_reshaped,
-            use_reentrant=False,
-        )  # (B*M, N_patches, embed_dim)
+        if self.use_checkpoint:
+            latent = checkpoint(
+                self.encoder,
+                input_data_reshaped,
+                daily_mask_reshaped,
+                use_reentrant=False,
+            )  # (B*M, N_patches, embed_dim)
+        else:
+            latent = self.encoder(
+                input_data_reshaped, daily_mask_reshaped
+            )  # (B*M, N_patches, embed_dim)
 
         # Step 2: Aggregate temporal information for each spatial patch
         # temporal input shape = (B, M*T*H*W, C),  C: embedding dimension
         # temporal output shape = (B, M, H*W, C)  C: embedding dimension
         embed_dim = latent.shape[-1]
         latent = latent.view(B, M, Tp, Hp, Wp, embed_dim)
 
-        agg_latent = checkpoint(
-            self.temporal, latent, M, daily_timef, padded_days_mask, use_reentrant=False
-        )  # (B, M, Hp*Wp, embed_dim)
+        if self.use_checkpoint:
+            agg_latent = checkpoint(
+                self.temporal,
+                latent,
+                M,
+                daily_timef,
+                padded_days_mask,
+                use_reentrant=False,
+            )  # (B, M, Hp*Wp, embed_dim)
+        else:
+            agg_latent = self.temporal(
+                latent, M, daily_timef, padded_days_mask
+            )  # (B, M, Hp*Wp, embed_dim)
 
         # Step 3: Add geo position and scale encodings
-        geo_emb = checkpoint(
-            self.geo_embedding,
-            geo_pos_embedding_patch,
-            scale_feature_patch,
-            use_reentrant=False,
-        )[:, None, None, :]  # (B,1,1,E)
+        if self.use_checkpoint:
+            geo_emb = checkpoint(
+                self.geo_embedding,
+                geo_pos_embedding_patch,
+                scale_feature_patch,
+                use_reentrant=False,
+            )[:, None, None, :]  # (B,1,1,E)
+        else:
+            geo_emb = self.geo_embedding(geo_pos_embedding_patch, scale_feature_patch)[
+                :, None, None, :
+            ]  # (B,1,1,E)
 
         # Broadcasting: same geo embedding for all M months at each Hp*Wp location
         # we use weighted mean patch embedding, see `geo_embedding_utils.py`
@@ -783,7 +806,10 @@ def forward(
         # Step 5: Decode to full-resolution 2D map
         # decoder input shape is (B, M*Hp*Wp, C), C: embedding dimension
         # decoder output shape is (B, M, H, W)
-        monthly_pred = checkpoint(
-            self.decoder, x, M, H, W, land_mask_patch, use_reentrant=False
-        )  # (B, M, H, W)
+        if self.use_checkpoint:
+            monthly_pred = checkpoint(
+                self.decoder, x, M, H, W, land_mask_patch, use_reentrant=False
+            )  # (B, M, H, W)
+        else:
+            monthly_pred = self.decoder(x, M, H, W, land_mask_patch)  # (B, M, H, W)
         return monthly_pred

climanet/utils.py

@@ -252,6 +252,7 @@ def compute_masked_loss(
 def save_model(model: torch.nn.Module, run_dir: str, verbose: bool) -> None:
     """Save model state and config to disk."""
     model_path = Path(run_dir) / "best_model.pth"
+    model = model.module if hasattr(model, "module") else model
     torch.save(
         {"model_state_dict": model.state_dict(), "model_config": model.config},
         model_path,

tests/test_dataset.py

@@ -113,5 +113,7 @@ def test_time_feature_generation():
     )
 
     sample = dataset[0]
-    expected_time_feature=torch.tensor([np.float32(2*np.pi*6/365.24),np.float32(0.)])
-    assert torch.equal(sample["daily_timef_patch"][0,5,:], expected_time_feature)
+    expected_time_feature = torch.tensor(
+        [np.float32(2 * np.pi * 6 / 365.24), np.float32(0.0)]
+    )
+    assert torch.equal(sample["daily_timef_patch"][0, 5, :], expected_time_feature)

tests/test_train.py

@@ -0,0 +1,99 @@
+import pytest
+import torch
+
+from climanet.train import _run_one_batch
+from climanet.st_encoder_decoder import SpatioTemporalModel
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+
+@pytest.fixture
+def create_dummy_batch():
+    # a dummy batch for testing
+    batch = {
+        "daily_patch": torch.rand(1, 1, 2, 31, 40, 40),
+        "monthly_patch": torch.rand(1, 2, 40, 40),
+        "daily_mask_patch": torch.rand(1, 1, 2, 31, 40, 40) > 0.5,  # boolean mask
+        "land_mask_patch": torch.rand(1, 40, 40) > 0.5,  # boolean mask
+        "daily_timef_patch": torch.rand(1, 2, 31, 2),
+        "padded_days_mask": torch.rand(1, 2, 31) > 0.5,  # boolean mask
+        "scale_feature_patch": torch.rand(1, 10),
+        "geo_pos_embedding_patch": torch.rand(1, 96),
+        "sh_embed_dim": torch.rand(1),
+        "harmonic_order": torch.rand(1),
+        "scale_f_dim": torch.rand(1),
+        "coords": torch.rand(1, 2),
+        "lat_patch": torch.rand(1, 40),
+        "lon_patch": torch.rand(1, 40),
+    }
+    return batch
+
+
+def test_model_meta_device(create_dummy_batch):
+    """Test that the model can run on a meta device and compute loss without errors.
+
+    Device is set to 'meta' for fast model construction, shape propagation,
+    and validating model architecture without executing ops.
+
+    """
+    batch = create_dummy_batch
+    model = SpatioTemporalModel(
+        patch_size=(1, 4, 4),
+        overlap=2,
+        num_months=2,
+        embed_dim=64,
+        dropout=0.2,
+        hidden=64,
+        use_checkpoint=True,
+    )
+    device = "meta"
+
+    model = model.to(device)
+    decoder = model.decoder
+    mean, std = torch.rand(1), torch.rand(1)
+    with torch.no_grad():
+        decoder.bias.copy_(mean)
+        decoder.scale.copy_(std + 1e-6)
+
+    batch = {k: v.to(device, non_blocking=False) for k, v in batch.items()}
+
+    model.train()
+    loss = _run_one_batch(model, batch)
+    loss.backward()
+
+    assert loss.device.type == "meta"
+
+
+def test_model_fake_tensor(create_dummy_batch):
+    """Test that the model can run with fake tensors and compute loss without errors.
+
+    This test uses fake tensors to test operator dispatch, device placement, and
+    graph correctness without executing real kernels.
+
+    For this test, checkpointing is disabled to avoid issues with fake tensors and autograd.
+
+    """
+    batch = create_dummy_batch
+    model = SpatioTemporalModel(
+        patch_size=(1, 4, 4),
+        overlap=2,
+        num_months=2,
+        embed_dim=64,
+        dropout=0.2,
+        hidden=64,
+        use_checkpoint=False,
+    )
+    device = "cpu"
+
+    model = model.to(device)
+    decoder = model.decoder
+    mean, std = torch.rand(1), torch.rand(1)
+    with torch.no_grad():
+        decoder.bias.copy_(mean)
+        decoder.scale.copy_(std + 1e-6)
+
+    with FakeTensorMode(allow_non_fake_inputs=True) as mode:
+        batch = {k: mode.from_tensor(v) for k, v in batch.items()}
+        loss = _run_one_batch(model, batch)
+        loss.backward()
+
+    assert loss.device.type == "cpu"