prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  [ShardedGradScaler] Handle optimizers not being OSS instances
  ## 🚀 Feature
  Some frameworks (Classy for instance) change the optimizers to enable different non-pytorch-compliant features. ShardedGradScaler asserts on that
  
  ## Motivation
  Enable ShardedDDP + AMP on Classy-like frameworks
  
  ## Pitch
  Remove the assert, replace by a one time warning
  
  ## Alternatives
  Not doing anything
  
  ## Additional context
  @mannatsingh 
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `fairscale/optim/grad_scaler.py`
  Content:
  ```
  1 # Copyright (c) Facebook, Inc. and its affiliates.
  2 #
  3 # This source code is licensed under the MIT license found in the
  4 # LICENSE file in the root directory of this source tree.
  5 
  6 from typing import Dict
  7 
  8 import torch
  9 from torch.cuda.amp import GradScaler as TorchGradScaler
  10 import torch.distributed as dist
  11 from torch.optim import Optimizer
  12 
  13 from .oss import OSS
  14 
  15 
  16 class GradScaler(TorchGradScaler):
  17     def _unscale_grads_(
  18         self, optimizer: Optimizer, inv_scale: torch.Tensor, found_inf: torch.Tensor, allow_fp16: bool
  19     ) -> Dict[torch.device, torch.Tensor]:
  20         return super()._unscale_grads_(optimizer, inv_scale, found_inf, True)
  21 
  22 
  23 class ShardedGradScaler(TorchGradScaler):
  24     """
  25     A shard-aware :class:`GradScaler<torch.cuda.amp.GradScaler>`, to be used in conjunction with
  26     :class:`OSS` and :class:`ShardedOptimizer`.
  27 
  28     Interface and usecases are not changed, more explanations can be found in the corresponding pytorch
  29     documentation https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler
  30     """
  31 
  32     def __init__(self) -> None:
  33         super().__init__()
  34 
  35     def unscale_(self, optimizer: Optimizer) -> None:
  36         assert isinstance(optimizer, OSS), "ShardedGradScaler is to be used in combination with a sharded optimizer"
  37 
  38         # Call the upstream unscale_ method which will only act on this rank's gradients
  39         super().unscale_(optimizer)
  40 
  41         # Synchronize the detected inf across the ranks
  42         optimizer_state = self._per_optimizer_states[id(optimizer)]
  43         handles = [dist.all_reduce(v, async_op=True) for v in optimizer_state["found_inf_per_device"].values()]
  44 
  45         # Make sure that the calls are done before moving out
  46         _ = list(map(lambda x: x.wait(), handles))
  47
  ```
  --- 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/fairscale/optim/grad_scaler.py b/fairscale/optim/grad_scaler.py
  --- a/fairscale/optim/grad_scaler.py
  +++ b/fairscale/optim/grad_scaler.py
  @@ -3,6 +3,7 @@
   # This source code is licensed under the MIT license found in the
   # LICENSE file in the root directory of this source tree.
   
  +import logging
   from typing import Dict
   
   import torch
  @@ -31,9 +32,16 @@
   
       def __init__(self) -> None:
           super().__init__()
  +        self.display_warning = True
   
       def unscale_(self, optimizer: Optimizer) -> None:
  -        assert isinstance(optimizer, OSS), "ShardedGradScaler is to be used in combination with a sharded optimizer"
  +        # Could be a mistake, this scaler is supposed to work with ZeroRedundancyOptimizer only
  +        if self.display_warning and not isinstance(optimizer, OSS):
  +            logging.warning(
  +                "ShardedGradScaler is to be used in combination with a sharded optimizer, this could not be checked"
  +            )
  +
  +        self.display_warning = False  # Only warn once
   
           # Call the upstream unscale_ method which will only act on this rank's gradients
           super().unscale_(optimizer)
  
verification_info: {"golden_diff": "diff --git a/fairscale/optim/grad_scaler.py b/fairscale/optim/grad_scaler.py\n--- a/fairscale/optim/grad_scaler.py\n+++ b/fairscale/optim/grad_scaler.py\n@@ -3,6 +3,7 @@\n # This source code is licensed under the MIT license found in the\n # LICENSE file in the root directory of this source tree.\n \n+import logging\n from typing import Dict\n \n import torch\n@@ -31,9 +32,16 @@\n \n     def __init__(self) -> None:\n         super().__init__()\n+        self.display_warning = True\n \n     def unscale_(self, optimizer: Optimizer) -> None:\n-        assert isinstance(optimizer, OSS), \"ShardedGradScaler is to be used in combination with a sharded optimizer\"\n+        # Could be a mistake, this scaler is supposed to work with ZeroRedundancyOptimizer only\n+        if self.display_warning and not isinstance(optimizer, OSS):\n+            logging.warning(\n+                \"ShardedGradScaler is to be used in combination with a sharded optimizer, this could not be checked\"\n+            )\n+\n+        self.display_warning = False  # Only warn once\n \n         # Call the upstream unscale_ method which will only act on this rank's gradients\n         super().unscale_(optimizer)\n", "issue": "[ShardedGradScaler] Handle optimizers not being OSS instances\n## \ud83d\ude80 Feature\r\nSome frameworks (Classy for instance) change the optimizers to enable different non-pytorch-compliant features. ShardedGradScaler asserts on that\r\n\r\n## Motivation\r\nEnable ShardedDDP + AMP on Classy-like frameworks\r\n\r\n## Pitch\r\nRemove the assert, replace by a one time warning\r\n\r\n## Alternatives\r\nNot doing anything\r\n\r\n## Additional context\r\n@mannatsingh \n", "before_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates.\n#\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\n\nfrom typing import Dict\n\nimport torch\nfrom torch.cuda.amp import GradScaler as TorchGradScaler\nimport torch.distributed as dist\nfrom torch.optim import Optimizer\n\nfrom .oss import OSS\n\n\nclass GradScaler(TorchGradScaler):\n    def _unscale_grads_(\n        self, optimizer: Optimizer, inv_scale: torch.Tensor, found_inf: torch.Tensor, allow_fp16: bool\n    ) -> Dict[torch.device, torch.Tensor]:\n        return super()._unscale_grads_(optimizer, inv_scale, found_inf, True)\n\n\nclass ShardedGradScaler(TorchGradScaler):\n    \"\"\"\n    A shard-aware :class:`GradScaler<torch.cuda.amp.GradScaler>`, to be used in conjunction with\n    :class:`OSS` and :class:`ShardedOptimizer`.\n\n    Interface and usecases are not changed, more explanations can be found in the corresponding pytorch\n    documentation https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler\n    \"\"\"\n\n    def __init__(self) -> None:\n        super().__init__()\n\n    def unscale_(self, optimizer: Optimizer) -> None:\n        assert isinstance(optimizer, OSS), \"ShardedGradScaler is to be used in combination with a sharded optimizer\"\n\n        # Call the upstream unscale_ method which will only act on this rank's gradients\n        super().unscale_(optimizer)\n\n        # Synchronize the detected inf across the ranks\n        optimizer_state = self._per_optimizer_states[id(optimizer)]\n        handles = [dist.all_reduce(v, async_op=True) for v in optimizer_state[\"found_inf_per_device\"].values()]\n\n        # Make sure that the calls are done before moving out\n        _ = list(map(lambda x: x.wait(), handles))\n", "path": "fairscale/optim/grad_scaler.py"}], "after_files": [{"content": "# Copyright (c) Facebook, Inc. and its affiliates.\n#\n# This source code is licensed under the MIT license found in the\n# LICENSE file in the root directory of this source tree.\n\nimport logging\nfrom typing import Dict\n\nimport torch\nfrom torch.cuda.amp import GradScaler as TorchGradScaler\nimport torch.distributed as dist\nfrom torch.optim import Optimizer\n\nfrom .oss import OSS\n\n\nclass GradScaler(TorchGradScaler):\n    def _unscale_grads_(\n        self, optimizer: Optimizer, inv_scale: torch.Tensor, found_inf: torch.Tensor, allow_fp16: bool\n    ) -> Dict[torch.device, torch.Tensor]:\n        return super()._unscale_grads_(optimizer, inv_scale, found_inf, True)\n\n\nclass ShardedGradScaler(TorchGradScaler):\n    \"\"\"\n    A shard-aware :class:`GradScaler<torch.cuda.amp.GradScaler>`, to be used in conjunction with\n    :class:`OSS` and :class:`ShardedOptimizer`.\n\n    Interface and usecases are not changed, more explanations can be found in the corresponding pytorch\n    documentation https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler\n    \"\"\"\n\n    def __init__(self) -> None:\n        super().__init__()\n        self.display_warning = True\n\n    def unscale_(self, optimizer: Optimizer) -> None:\n        # Could be a mistake, this scaler is supposed to work with ZeroRedundancyOptimizer only\n        if self.display_warning and not isinstance(optimizer, OSS):\n            logging.warning(\n                \"ShardedGradScaler is to be used in combination with a sharded optimizer, this could not be checked\"\n            )\n\n        self.display_warning = False  # Only warn once\n\n        # Call the upstream unscale_ method which will only act on this rank's gradients\n        super().unscale_(optimizer)\n\n        # Synchronize the detected inf across the ranks\n        optimizer_state = self._per_optimizer_states[id(optimizer)]\n        handles = [dist.all_reduce(v, async_op=True) for v in optimizer_state[\"found_inf_per_device\"].values()]\n\n        # Make sure that the calls are done before moving out\n        _ = list(map(lambda x: x.wait(), handles))\n", "path": "fairscale/optim/grad_scaler.py"}]}