prompt: You will be provided with a partial code base and an issue statement explaining a problem to resolve.
  <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
  
  </issue>
  <code>
  [start of pytorch_lightning/loggers/wandb.py]
  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 of pytorch_lightning/loggers/wandb.py]
  </code>
  
  I need you to solve this issue by generating a single patch file that I can apply directly to this repository using git apply. Please respond with a single patch file in the following format.
  <patch>
  diff --git a/file.py b/file.py
  --- a/file.py
  +++ b/file.py
  @@ -1,27 +1,35 @@
   def euclidean(a, b):
  -    while b:
  -        a, b = b, a % b
  -    return a
  +    if b == 0:
  +        return a
  +    return euclidean(b, a % b)
   
   
   def bresenham(x0, y0, x1, y1):
       points = []
       dx = abs(x1 - x0)
       dy = abs(y1 - y0)
  -    sx = 1 if x0 < x1 else -1
  -    sy = 1 if y0 < y1 else -1
  -    err = dx - dy
  +    x, y = x0, y0
  +    sx = -1 if x0 > x1 else 1
  +    sy = -1 if y0 > y1 else 1
   
  -    while True:
  -        points.append((x0, y0))
  -        if x0 == x1 and y0 == y1:
  -            break
  -        e2 = 2 * err
  -        if e2 > -dy:
  -            err -= dy
  -            x0 += sx
  -        if e2 < dx:
  -            err += dx
  -            y0 += sy
  +    if dx > dy:
  +        err = dx / 2.0
  +        while x != x1:
  +            points.append((x, y))
  +            err -= dy
  +            if err < 0:
  +                y += sy
  +                err += dx
  +            x += sx
  +    else:
  +        err = dy / 2.0
  +        while y != y1:
  +            points.append((x, y))
  +            err -= dx
  +            if err < 0:
  +                x += sx
  +                err += dy
  +            y += sy
  +
  +    points.append((x, y))
       return points
  
  </patch>
golden_diff: 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()
  
verification_info: {"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"}]}