[bug] fixed grad scaler compatibility with torch 1.8
diff --git a/colossalai/amp/torch_amp/_grad_scaler.py b/colossalai/amp/torch_amp/_grad_scaler.py
index 48c7eb949b0b..d5e526d0c00a 100644
--- a/colossalai/amp/torch_amp/_grad_scaler.py
+++ b/colossalai/amp/torch_amp/_grad_scaler.py
@@ -12,6 +12,7 @@
 import torch.distributed as dist
 from colossalai.core import global_context as gpc
 from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+from packaging import version
 
 
 class _MultiDeviceReplicator(object):
@@ -122,6 +123,14 @@ def __init__(self, init_scale=2.**16, growth_factor=2.0, backoff_factor=0.5, gro
         else:
             self._enabled = enabled
 
+        # check version
+        torch_version = version.parse(torch.__version__)
+        assert torch_version.major == 1
+        if torch_version.minor > 8:
+            self._higher_than_torch18 = True
+        else:
+            self._higher_than_torch18 = False
+
         if self._enabled:
             assert growth_factor > 1.0, "The growth factor must be > 1.0."
             assert backoff_factor < 1.0, "The backoff factor must be < 1.0."
@@ -404,8 +413,12 @@ def update(self, new_scale=None):
                 for i in range(1, len(found_infs)):
                     found_inf_combined += found_infs[i]
 
-            torch._amp_update_scale_(_scale, _growth_tracker, found_inf_combined, self._growth_factor,
-                                     self._backoff_factor, self._growth_interval)
+            if self._higher_than_torch18:
+                torch._amp_update_scale_(_scale, _growth_tracker, found_inf_combined, self._growth_factor,
+                                         self._backoff_factor, self._growth_interval)
+            else:
+                self._scale = torch._amp_update_scale(_growth_tracker, _scale, found_inf_combined, self._growth_factor,
+                                                      self._backoff_factor, self._growth_interval)
 
         # To prepare for next iteration, clear the data collected from optimizers this iteration.
         self._per_optimizer_states = defaultdict(_refresh_per_optimizer_state)

This fixed #734 as part of the effort to tackle #733 .

The screenshot of the successful unit test is shown below.
<img width="1913" alt="Screenshot 2022-04-12 at 3 18 47 PM" src="https://user-images.githubusercontent.com/31818963/162903213-633f92fc-e532-4cac-a9b7-0855313ff283.png">

