We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Wandb logger doesn't upload saved model checkpoint for final epoch
## 🐛 Bug

When training a model on the TPU and using the wandb logger, the checkpoint for the last epoch trained doesn't get uploaded to wandb.

### To Reproduce

Colab notebook: https://colab.research.google.com/drive/1oPaRWGZcz6YEol012xFADN42LV-jowtT

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pytorch_lightning/loggers/wandb.py`
Content:
```
1 r"""
2 
3 .. _wandb:
4 
5 WandbLogger
6 -------------
7 """
8 import os
9 from argparse import Namespace
10 from typing import Optional, List, Dict, Union, Any
11 
12 import torch.nn as nn
13 
14 try:
15     import wandb
16     from wandb.wandb_run import Run
17 except ImportError:  # pragma: no-cover
18     raise ImportError('You want to use `wandb` logger which is not installed yet,'  # pragma: no-cover
19                       ' install it with `pip install wandb`.')
20 
21 from pytorch_lightning.loggers.base import LightningLoggerBase, rank_zero_only
22 
23 
24 class WandbLogger(LightningLoggerBase):
25     """
26     Logger for `W&B <https://www.wandb.com/>`_.
27 
28     Args:
29         name (str): display name for the run.
30         save_dir (str): path where data is saved.
31         offline (bool): run offline (data can be streamed later to wandb servers).
32         id or version (str): sets the version, mainly used to resume a previous run.
33         anonymous (bool): enables or explicitly disables anonymous logging.
34         project (str): the name of the project to which this run will belong.
35         tags (list of str): tags associated with this run.
36 
37     Example
38     --------
39     .. code-block:: python
40 
41         from pytorch_lightning.loggers import WandbLogger
42         from pytorch_lightning import Trainer
43 
44         wandb_logger = WandbLogger()
45         trainer = Trainer(logger=wandb_logger)
46     """
47 
48     def __init__(self, name: Optional[str] = None, save_dir: Optional[str] = None,
49                  offline: bool = False, id: Optional[str] = None, anonymous: bool = False,
50                  version: Optional[str] = None, project: Optional[str] = None,
51                  tags: Optional[List[str]] = None, experiment=None, entity=None):
52         super().__init__()
53         self._name = name
54         self._save_dir = save_dir
55         self._anonymous = 'allow' if anonymous else None
56         self._id = version or id
57         self._tags = tags
58         self._project = project
59         self._experiment = experiment
60         self._offline = offline
61         self._entity = entity
62 
63     def __getstate__(self):
64         state = self.__dict__.copy()
65         # cannot be pickled
66         state['_experiment'] = None
67         # args needed to reload correct experiment
68         state['_id'] = self.experiment.id
69         return state
70 
71     @property
72     def experiment(self) -> Run:
73         r"""
74 
75           Actual wandb object. To use wandb features do the following.
76 
77           Example::
78 
79               self.logger.experiment.some_wandb_function()
80 
81           """
82         if self._experiment is None:
83             if self._offline:
84                 os.environ['WANDB_MODE'] = 'dryrun'
85             self._experiment = wandb.init(
86                 name=self._name, dir=self._save_dir, project=self._project, anonymous=self._anonymous,
87                 id=self._id, resume='allow', tags=self._tags, entity=self._entity)
88         return self._experiment
89 
90     def watch(self, model: nn.Module, log: str = 'gradients', log_freq: int = 100):
91         wandb.watch(model, log=log, log_freq=log_freq)
92 
93     @rank_zero_only
94     def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
95         params = self._convert_params(params)
96         self.experiment.config.update(params)
97 
98     @rank_zero_only
99     def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:
100         if step is not None:
101             metrics['global_step'] = step
102         self.experiment.log(metrics)
103 
104     @rank_zero_only
105     def finalize(self, status: str = 'success') -> None:
106         try:
107             exit_code = 0 if status == 'success' else 1
108             wandb.join(exit_code)
109         except TypeError:
110             wandb.join()
111 
112     @property
113     def name(self) -> str:
114         return self.experiment.project_name()
115 
116     @property
117     def version(self) -> str:
118         return self.experiment.id
119
```
--- 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/pytorch_lightning/loggers/wandb.py b/pytorch_lightning/loggers/wandb.py
--- a/pytorch_lightning/loggers/wandb.py
+++ b/pytorch_lightning/loggers/wandb.py
@@ -101,14 +101,6 @@
             metrics['global_step'] = step
         self.experiment.log(metrics)
 
-    @rank_zero_only
-    def finalize(self, status: str = 'success') -> None:
-        try:
-            exit_code = 0 if status == 'success' else 1
-            wandb.join(exit_code)
-        except TypeError:
-            wandb.join()
-
     @property
     def name(self) -> str:
         return self.experiment.project_name()

{"golden_diff": "diff --git a/pytorch_lightning/loggers/wandb.py b/pytorch_lightning/loggers/wandb.py\n--- a/pytorch_lightning/loggers/wandb.py\n+++ b/pytorch_lightning/loggers/wandb.py\n@@ -101,14 +101,6 @@\n             metrics['global_step'] = step\n         self.experiment.log(metrics)\n \n-    @rank_zero_only\n-    def finalize(self, status: str = 'success') -> None:\n-        try:\n-            exit_code = 0 if status == 'success' else 1\n-            wandb.join(exit_code)\n-        except TypeError:\n-            wandb.join()\n-\n     @property\n     def name(self) -> str:\n         return self.experiment.project_name()\n", "issue": "Wandb logger doesn't upload saved model checkpoint for final epoch\n## \ud83d\udc1b Bug\r\n\r\nWhen training a model on the TPU and using the wandb logger, the checkpoint for the last epoch trained doesn't get uploaded to wandb.\r\n\r\n### To Reproduce\r\n\r\nColab notebook: https://colab.research.google.com/drive/1oPaRWGZcz6YEol012xFADN42LV-jowtT\n", "before_files": [{"content": "r\"\"\"\n\n.. _wandb:\n\nWandbLogger\n-------------\n\"\"\"\nimport os\nfrom argparse import Namespace\nfrom typing import Optional, List, Dict, Union, Any\n\nimport torch.nn as nn\n\ntry:\n    import wandb\n    from wandb.wandb_run import Run\nexcept ImportError:  # pragma: no-cover\n    raise ImportError('You want to use `wandb` logger which is not installed yet,'  # pragma: no-cover\n                      ' install it with `pip install wandb`.')\n\nfrom pytorch_lightning.loggers.base import LightningLoggerBase, rank_zero_only\n\n\nclass WandbLogger(LightningLoggerBase):\n    \"\"\"\n    Logger for `W&B <https://www.wandb.com/>`_.\n\n    Args:\n        name (str): display name for the run.\n        save_dir (str): path where data is saved.\n        offline (bool): run offline (data can be streamed later to wandb servers).\n        id or version (str): sets the version, mainly used to resume a previous run.\n        anonymous (bool): enables or explicitly disables anonymous logging.\n        project (str): the name of the project to which this run will belong.\n        tags (list of str): tags associated with this run.\n\n    Example\n    --------\n    .. code-block:: python\n\n        from pytorch_lightning.loggers import WandbLogger\n        from pytorch_lightning import Trainer\n\n        wandb_logger = WandbLogger()\n        trainer = Trainer(logger=wandb_logger)\n    \"\"\"\n\n    def __init__(self, name: Optional[str] = None, save_dir: Optional[str] = None,\n                 offline: bool = False, id: Optional[str] = None, anonymous: bool = False,\n                 version: Optional[str] = None, project: Optional[str] = None,\n                 tags: Optional[List[str]] = None, experiment=None, entity=None):\n        super().__init__()\n        self._name = name\n        self._save_dir = save_dir\n        self._anonymous = 'allow' if anonymous else None\n        self._id = version or id\n        self._tags = tags\n        self._project = project\n        self._experiment = experiment\n        self._offline = offline\n        self._entity = entity\n\n    def __getstate__(self):\n        state = self.__dict__.copy()\n        # cannot be pickled\n        state['_experiment'] = None\n        # args needed to reload correct experiment\n        state['_id'] = self.experiment.id\n        return state\n\n    @property\n    def experiment(self) -> Run:\n        r\"\"\"\n\n          Actual wandb object. To use wandb features do the following.\n\n          Example::\n\n              self.logger.experiment.some_wandb_function()\n\n          \"\"\"\n        if self._experiment is None:\n            if self._offline:\n                os.environ['WANDB_MODE'] = 'dryrun'\n            self._experiment = wandb.init(\n                name=self._name, dir=self._save_dir, project=self._project, anonymous=self._anonymous,\n                id=self._id, resume='allow', tags=self._tags, entity=self._entity)\n        return self._experiment\n\n    def watch(self, model: nn.Module, log: str = 'gradients', log_freq: int = 100):\n        wandb.watch(model, log=log, log_freq=log_freq)\n\n    @rank_zero_only\n    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:\n        params = self._convert_params(params)\n        self.experiment.config.update(params)\n\n    @rank_zero_only\n    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:\n        if step is not None:\n            metrics['global_step'] = step\n        self.experiment.log(metrics)\n\n    @rank_zero_only\n    def finalize(self, status: str = 'success') -> None:\n        try:\n            exit_code = 0 if status == 'success' else 1\n            wandb.join(exit_code)\n        except TypeError:\n            wandb.join()\n\n    @property\n    def name(self) -> str:\n        return self.experiment.project_name()\n\n    @property\n    def version(self) -> str:\n        return self.experiment.id\n", "path": "pytorch_lightning/loggers/wandb.py"}], "after_files": [{"content": "r\"\"\"\n\n.. _wandb:\n\nWandbLogger\n-------------\n\"\"\"\nimport os\nfrom argparse import Namespace\nfrom typing import Optional, List, Dict, Union, Any\n\nimport torch.nn as nn\n\ntry:\n    import wandb\n    from wandb.wandb_run import Run\nexcept ImportError:  # pragma: no-cover\n    raise ImportError('You want to use `wandb` logger which is not installed yet,'  # pragma: no-cover\n                      ' install it with `pip install wandb`.')\n\nfrom pytorch_lightning.loggers.base import LightningLoggerBase, rank_zero_only\n\n\nclass WandbLogger(LightningLoggerBase):\n    \"\"\"\n    Logger for `W&B <https://www.wandb.com/>`_.\n\n    Args:\n        name (str): display name for the run.\n        save_dir (str): path where data is saved.\n        offline (bool): run offline (data can be streamed later to wandb servers).\n        id or version (str): sets the version, mainly used to resume a previous run.\n        anonymous (bool): enables or explicitly disables anonymous logging.\n        project (str): the name of the project to which this run will belong.\n        tags (list of str): tags associated with this run.\n\n    Example\n    --------\n    .. code-block:: python\n\n        from pytorch_lightning.loggers import WandbLogger\n        from pytorch_lightning import Trainer\n\n        wandb_logger = WandbLogger()\n        trainer = Trainer(logger=wandb_logger)\n    \"\"\"\n\n    def __init__(self, name: Optional[str] = None, save_dir: Optional[str] = None,\n                 offline: bool = False, id: Optional[str] = None, anonymous: bool = False,\n                 version: Optional[str] = None, project: Optional[str] = None,\n                 tags: Optional[List[str]] = None, experiment=None, entity=None):\n        super().__init__()\n        self._name = name\n        self._save_dir = save_dir\n        self._anonymous = 'allow' if anonymous else None\n        self._id = version or id\n        self._tags = tags\n        self._project = project\n        self._experiment = experiment\n        self._offline = offline\n        self._entity = entity\n\n    def __getstate__(self):\n        state = self.__dict__.copy()\n        # cannot be pickled\n        state['_experiment'] = None\n        # args needed to reload correct experiment\n        state['_id'] = self.experiment.id\n        return state\n\n    @property\n    def experiment(self) -> Run:\n        r\"\"\"\n\n          Actual wandb object. To use wandb features do the following.\n\n          Example::\n\n              self.logger.experiment.some_wandb_function()\n\n          \"\"\"\n        if self._experiment is None:\n            if self._offline:\n                os.environ['WANDB_MODE'] = 'dryrun'\n            self._experiment = wandb.init(\n                name=self._name, dir=self._save_dir, project=self._project, anonymous=self._anonymous,\n                id=self._id, resume='allow', tags=self._tags, entity=self._entity)\n        return self._experiment\n\n    def watch(self, model: nn.Module, log: str = 'gradients', log_freq: int = 100):\n        wandb.watch(model, log=log, log_freq=log_freq)\n\n    @rank_zero_only\n    def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:\n        params = self._convert_params(params)\n        self.experiment.config.update(params)\n\n    @rank_zero_only\n    def log_metrics(self, metrics: Dict[str, float], step: Optional[int] = None) -> None:\n        if step is not None:\n            metrics['global_step'] = step\n        self.experiment.log(metrics)\n\n    @property\n    def name(self) -> str:\n        return self.experiment.project_name()\n\n    @property\n    def version(self) -> str:\n        return self.experiment.id\n", "path": "pytorch_lightning/loggers/wandb.py"}]}