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()

```

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."""

{"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"}]}