prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Checkpoint saving not working as expected
  ### Description
  
  After migrating to release 0.5.0 noticed that checkpoint saving is not working as expected.
  
  ## description
  tried different configuration e.g., `checkpoint_callback = ModelCheckpoint(monitor="val_loss", dirpath=ckpt_dir, save_last=True, every_n_epochs=1, save_top_k=1)` for example when running 20-30 epochs for training a model.
  after training was completed could not find the ckpt file. what was found was a single ckpt file of the first epoch only, in a wrong directory.
  
  ## severance
  the bug is very limiting. for example, after hours of training a model, there is no way to load the model from a checkpoint to run inference. the single shot to run inference was during the same run. 
  
  ## expected behavior
  using a given configuration expected to see:
  - checkpoint files saved every number of epoch
  - the last epoch checkpoint file
  - the checkpoints should have been saved to the given directory
  
  ## observed behavior
  - after training several epochs only the first was saved.
  - the single checkpoint were saved to another directory under the logger output
  
  ## initial investigation
  1. checkpoint callback created and training fit called
  2. later, see image and call stack:
  seems like c'tor called again with save_last=None
  ![image](https://github.com/microsoft/torchgeo/assets/11072478/667b6cc1-8b7b-4ede-a61d-c365e851ab59)
  
  3. when saving later supposed to happen, the save_last is None:
  ![image](https://github.com/microsoft/torchgeo/assets/11072478/2dd68eff-3c43-4ccc-8ce7-4a463044000b)
  
  4. last checkpoint saving is skipped
  
  
  ### Steps to reproduce
  
  1. create a checkpoint callback and use different checkpoints saving parameters e.g., checkpoint_callback = ModelCheckpoint(monitor="val_loss", dirpath=ckpt_dir, save_last=True, every_n_epochs=1, save_top_k=1)
  2. call trainer fit and run several epochs
  3. check expected results:
  - saving location as expected e.g., under C:\foo 
  - check last epoch checkpoint saved - must have last.ckpt
  - check how many checkpoints were saved e.g., every 2 etc
  
  ### Version
  
  torchgeo version 0.5.0, lightning version 2.0.9
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `torchgeo/trainers/base.py`
  Content:
  ```
  1 # Copyright (c) Microsoft Corporation. All rights reserved.
  2 # Licensed under the MIT License.
  3 
  4 """Base classes for all :mod:`torchgeo` trainers."""
  5 
  6 from abc import ABC, abstractmethod
  7 from typing import Any
  8 
  9 import lightning
  10 from lightning.pytorch import LightningModule
  11 from lightning.pytorch.callbacks import Callback, EarlyStopping, ModelCheckpoint
  12 from torch.optim import AdamW
  13 from torch.optim.lr_scheduler import ReduceLROnPlateau
  14 
  15 
  16 class BaseTask(LightningModule, ABC):
  17     """Abstract base class for all TorchGeo trainers.
  18 
  19     .. versionadded:: 0.5
  20     """
  21 
  22     #: Model to train.
  23     model: Any
  24 
  25     #: Performance metric to monitor in learning rate scheduler and callbacks.
  26     monitor = "val_loss"
  27 
  28     #: Whether the goal is to minimize or maximize the performance metric to monitor.
  29     mode = "min"
  30 
  31     def __init__(self) -> None:
  32         """Initialize a new BaseTask instance."""
  33         super().__init__()
  34         self.save_hyperparameters()
  35         self.configure_losses()
  36         self.configure_metrics()
  37         self.configure_models()
  38 
  39     def configure_callbacks(self) -> list[Callback]:
  40         """Initialize model-specific callbacks.
  41 
  42         Returns:
  43             List of callbacks to apply.
  44         """
  45         return [
  46             ModelCheckpoint(monitor=self.monitor, mode=self.mode),
  47             EarlyStopping(monitor=self.monitor, mode=self.mode),
  48         ]
  49 
  50     def configure_losses(self) -> None:
  51         """Initialize the loss criterion."""
  52 
  53     def configure_metrics(self) -> None:
  54         """Initialize the performance metrics."""
  55 
  56     @abstractmethod
  57     def configure_models(self) -> None:
  58         """Initialize the model."""
  59 
  60     def configure_optimizers(
  61         self,
  62     ) -> "lightning.pytorch.utilities.types.OptimizerLRSchedulerConfig":
  63         """Initialize the optimizer and learning rate scheduler.
  64 
  65         Returns:
  66             Optimizer and learning rate scheduler.
  67         """
  68         optimizer = AdamW(self.parameters(), lr=self.hparams["lr"])
  69         scheduler = ReduceLROnPlateau(optimizer, patience=self.hparams["patience"])
  70         return {
  71             "optimizer": optimizer,
  72             "lr_scheduler": {"scheduler": scheduler, "monitor": self.monitor},
  73         }
  74 
  75     def forward(self, *args: Any, **kwargs: Any) -> Any:
  76         """Forward pass of the model.
  77 
  78         Args:
  79             args: Arguments to pass to model.
  80             kwargs: Keyword arguments to pass to model.
  81 
  82         Returns:
  83             Output of the model.
  84         """
  85         return self.model(*args, **kwargs)
  86
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/torchgeo/trainers/base.py b/torchgeo/trainers/base.py
  --- a/torchgeo/trainers/base.py
  +++ b/torchgeo/trainers/base.py
  @@ -8,7 +8,6 @@
   
   import lightning
   from lightning.pytorch import LightningModule
  -from lightning.pytorch.callbacks import Callback, EarlyStopping, ModelCheckpoint
   from torch.optim import AdamW
   from torch.optim.lr_scheduler import ReduceLROnPlateau
   
  @@ -36,17 +35,6 @@
           self.configure_metrics()
           self.configure_models()
   
  -    def configure_callbacks(self) -> list[Callback]:
  -        """Initialize model-specific callbacks.
  -
  -        Returns:
  -            List of callbacks to apply.
  -        """
  -        return [
  -            ModelCheckpoint(monitor=self.monitor, mode=self.mode),
  -            EarlyStopping(monitor=self.monitor, mode=self.mode),
  -        ]
  -
       def configure_losses(self) -> None:
           """Initialize the loss criterion."""
  
verification_info: {"golden_diff": "diff --git a/torchgeo/trainers/base.py b/torchgeo/trainers/base.py\n--- a/torchgeo/trainers/base.py\n+++ b/torchgeo/trainers/base.py\n@@ -8,7 +8,6 @@\n \n import lightning\n from lightning.pytorch import LightningModule\n-from lightning.pytorch.callbacks import Callback, EarlyStopping, ModelCheckpoint\n from torch.optim import AdamW\n from torch.optim.lr_scheduler import ReduceLROnPlateau\n \n@@ -36,17 +35,6 @@\n         self.configure_metrics()\n         self.configure_models()\n \n-    def configure_callbacks(self) -> list[Callback]:\n-        \"\"\"Initialize model-specific callbacks.\n-\n-        Returns:\n-            List of callbacks to apply.\n-        \"\"\"\n-        return [\n-            ModelCheckpoint(monitor=self.monitor, mode=self.mode),\n-            EarlyStopping(monitor=self.monitor, mode=self.mode),\n-        ]\n-\n     def configure_losses(self) -> None:\n         \"\"\"Initialize the loss criterion.\"\"\"\n", "issue": "Checkpoint saving not working as expected\n### Description\r\n\r\nAfter migrating to release 0.5.0 noticed that checkpoint saving is not working as expected.\r\n\r\n## description\r\ntried different configuration e.g., `checkpoint_callback = ModelCheckpoint(monitor=\"val_loss\", dirpath=ckpt_dir, save_last=True, every_n_epochs=1, save_top_k=1)` for example when running 20-30 epochs for training a model.\r\nafter training was completed could not find the ckpt file. what was found was a single ckpt file of the first epoch only, in a wrong directory.\r\n\r\n## severance\r\nthe bug is very limiting. for example, after hours of training a model, there is no way to load the model from a checkpoint to run inference. the single shot to run inference was during the same run. \r\n\r\n## expected behavior\r\nusing a given configuration expected to see:\r\n- checkpoint files saved every number of epoch\r\n- the last epoch checkpoint file\r\n- the checkpoints should have been saved to the given directory\r\n\r\n## observed behavior\r\n- after training several epochs only the first was saved.\r\n- the single checkpoint were saved to another directory under the logger output\r\n\r\n## initial investigation\r\n1. checkpoint callback created and training fit called\r\n2. later, see image and call stack:\r\nseems like c'tor called again with save_last=None\r\n![image](https://github.com/microsoft/torchgeo/assets/11072478/667b6cc1-8b7b-4ede-a61d-c365e851ab59)\r\n\r\n3. when saving later supposed to happen, the save_last is None:\r\n![image](https://github.com/microsoft/torchgeo/assets/11072478/2dd68eff-3c43-4ccc-8ce7-4a463044000b)\r\n\r\n4. last checkpoint saving is skipped\r\n\r\n\r\n### Steps to reproduce\r\n\r\n1. create a checkpoint callback and use different checkpoints saving parameters e.g., checkpoint_callback = ModelCheckpoint(monitor=\"val_loss\", dirpath=ckpt_dir, save_last=True, every_n_epochs=1, save_top_k=1)\r\n2. call trainer fit and run several epochs\r\n3. check expected results:\r\n- saving location as expected e.g., under C:\\foo \r\n- check last epoch checkpoint saved - must have last.ckpt\r\n- check how many checkpoints were saved e.g., every 2 etc\r\n\r\n### Version\r\n\r\ntorchgeo version 0.5.0, lightning version 2.0.9\n", "before_files": [{"content": "# Copyright (c) Microsoft Corporation. All rights reserved.\n# Licensed under the MIT License.\n\n\"\"\"Base classes for all :mod:`torchgeo` trainers.\"\"\"\n\nfrom abc import ABC, abstractmethod\nfrom typing import Any\n\nimport lightning\nfrom lightning.pytorch import LightningModule\nfrom lightning.pytorch.callbacks import Callback, EarlyStopping, ModelCheckpoint\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\n\n\nclass BaseTask(LightningModule, ABC):\n    \"\"\"Abstract base class for all TorchGeo trainers.\n\n    .. versionadded:: 0.5\n    \"\"\"\n\n    #: Model to train.\n    model: Any\n\n    #: Performance metric to monitor in learning rate scheduler and callbacks.\n    monitor = \"val_loss\"\n\n    #: Whether the goal is to minimize or maximize the performance metric to monitor.\n    mode = \"min\"\n\n    def __init__(self) -> None:\n        \"\"\"Initialize a new BaseTask instance.\"\"\"\n        super().__init__()\n        self.save_hyperparameters()\n        self.configure_losses()\n        self.configure_metrics()\n        self.configure_models()\n\n    def configure_callbacks(self) -> list[Callback]:\n        \"\"\"Initialize model-specific callbacks.\n\n        Returns:\n            List of callbacks to apply.\n        \"\"\"\n        return [\n            ModelCheckpoint(monitor=self.monitor, mode=self.mode),\n            EarlyStopping(monitor=self.monitor, mode=self.mode),\n        ]\n\n    def configure_losses(self) -> None:\n        \"\"\"Initialize the loss criterion.\"\"\"\n\n    def configure_metrics(self) -> None:\n        \"\"\"Initialize the performance metrics.\"\"\"\n\n    @abstractmethod\n    def configure_models(self) -> None:\n        \"\"\"Initialize the model.\"\"\"\n\n    def configure_optimizers(\n        self,\n    ) -> \"lightning.pytorch.utilities.types.OptimizerLRSchedulerConfig\":\n        \"\"\"Initialize the optimizer and learning rate scheduler.\n\n        Returns:\n            Optimizer and learning rate scheduler.\n        \"\"\"\n        optimizer = AdamW(self.parameters(), lr=self.hparams[\"lr\"])\n        scheduler = ReduceLROnPlateau(optimizer, patience=self.hparams[\"patience\"])\n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": {\"scheduler\": scheduler, \"monitor\": self.monitor},\n        }\n\n    def forward(self, *args: Any, **kwargs: Any) -> Any:\n        \"\"\"Forward pass of the model.\n\n        Args:\n            args: Arguments to pass to model.\n            kwargs: Keyword arguments to pass to model.\n\n        Returns:\n            Output of the model.\n        \"\"\"\n        return self.model(*args, **kwargs)\n", "path": "torchgeo/trainers/base.py"}], "after_files": [{"content": "# Copyright (c) Microsoft Corporation. All rights reserved.\n# Licensed under the MIT License.\n\n\"\"\"Base classes for all :mod:`torchgeo` trainers.\"\"\"\n\nfrom abc import ABC, abstractmethod\nfrom typing import Any\n\nimport lightning\nfrom lightning.pytorch import LightningModule\nfrom torch.optim import AdamW\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\n\n\nclass BaseTask(LightningModule, ABC):\n    \"\"\"Abstract base class for all TorchGeo trainers.\n\n    .. versionadded:: 0.5\n    \"\"\"\n\n    #: Model to train.\n    model: Any\n\n    #: Performance metric to monitor in learning rate scheduler and callbacks.\n    monitor = \"val_loss\"\n\n    #: Whether the goal is to minimize or maximize the performance metric to monitor.\n    mode = \"min\"\n\n    def __init__(self) -> None:\n        \"\"\"Initialize a new BaseTask instance.\"\"\"\n        super().__init__()\n        self.save_hyperparameters()\n        self.configure_losses()\n        self.configure_metrics()\n        self.configure_models()\n\n    def configure_losses(self) -> None:\n        \"\"\"Initialize the loss criterion.\"\"\"\n\n    def configure_metrics(self) -> None:\n        \"\"\"Initialize the performance metrics.\"\"\"\n\n    @abstractmethod\n    def configure_models(self) -> None:\n        \"\"\"Initialize the model.\"\"\"\n\n    def configure_optimizers(\n        self,\n    ) -> \"lightning.pytorch.utilities.types.OptimizerLRSchedulerConfig\":\n        \"\"\"Initialize the optimizer and learning rate scheduler.\n\n        Returns:\n            Optimizer and learning rate scheduler.\n        \"\"\"\n        optimizer = AdamW(self.parameters(), lr=self.hparams[\"lr\"])\n        scheduler = ReduceLROnPlateau(optimizer, patience=self.hparams[\"patience\"])\n        return {\n            \"optimizer\": optimizer,\n            \"lr_scheduler\": {\"scheduler\": scheduler, \"monitor\": self.monitor},\n        }\n\n    def forward(self, *args: Any, **kwargs: Any) -> Any:\n        \"\"\"Forward pass of the model.\n\n        Args:\n            args: Arguments to pass to model.\n            kwargs: Keyword arguments to pass to model.\n\n        Returns:\n            Output of the model.\n        \"\"\"\n        return self.model(*args, **kwargs)\n", "path": "torchgeo/trainers/base.py"}]}