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

```

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)

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