Merge pull request #140 from allenai/favyen/upgrade-pytorch

Bump Lightning to 2.5.0.post0

extra_requirements.txt

@@ -3,7 +3,7 @@ accelerate>=1.0
 cdsapi>=0.7.4
 earthengine-api>=0.1
 einops>=0.8
-gcsfs==2024.12
+gcsfs==2025.2.0
 google-cloud-bigquery>=2.18
 google-cloud-storage>=2.18
 interrogate>=1.7
@@ -14,7 +14,7 @@ planetary_computer>=1.0
 pycocotools>=2.0
 pystac_client>=0.8
 rtree>=1.2
-s3fs==2024.12
+s3fs==2025.2.0
 satlaspretrain_models>=0.3
 scipy>=1.13
 transformers>=4.45

requirements.txt

@@ -1,14 +1,14 @@
 boto3>=1.35
 class_registry>=2.1
 fiona>=1.10
-fsspec==2024.12
-lightning[pytorch-extra]==2.5.0
+fsspec==2025.2.0
+lightning[pytorch-extra]==2.5.0.post0
 Pillow>=11.0
 pyproj>=3.7
 python-dateutil>=2.9
 pytimeparse>=1.1
 rasterio>=1.4
 shapely>=2.0
-torchvision==0.20.0
+torchvision==0.21.0
 tqdm>=4.66
 universal_pathlib>=0.2.5

rslearn/train/lightning_module.py

@@ -5,7 +5,7 @@
 
 import lightning as L
 import torch
-from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfigType
+from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfig
 from PIL import Image
 from torch.optim import AdamW
 from torch.optim.lr_scheduler import ReduceLROnPlateau
@@ -154,7 +154,7 @@ def on_fit_start(self) -> None:
                     f"warning: restore yielded missing_keys={missing_keys} and unexpected_keys={unexpected_keys}"
                 )
 
-    def configure_optimizers(self) -> OptimizerLRSchedulerConfigType:
+    def configure_optimizers(self) -> OptimizerLRSchedulerConfig:
         """Initialize the optimizer and learning rate scheduler.
 
         Returns:

tests/integration/train/test_freeze_unfreeze.py

@@ -2,7 +2,7 @@
 
 import lightning.pytorch as pl
 import pytest
-from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfigType
+from lightning.pytorch.utilities.types import OptimizerLRSchedulerConfig
 
 from rslearn.dataset import Dataset
 from rslearn.models.pooling_decoder import PoolingDecoder
@@ -36,7 +36,7 @@ class LMWithCustomPlateau(RslearnLightningModule):
     epoch.
     """
 
-    def configure_optimizers(self) -> OptimizerLRSchedulerConfigType:
+    def configure_optimizers(self) -> OptimizerLRSchedulerConfig:
         """Initialize the optimizer and learning rate scheduler.
 
         Returns: