[hotfix] fix initialize bug with zero (#442)

colossalai/amp/__init__.py

@@ -11,17 +11,13 @@ from .apex_amp import convert_to_apex_amp
 from .naive_amp import convert_to_naive_amp
 
 
-def convert_to_amp(model: nn.Module,
-                   optimizer: Optimizer,
-                   criterion: _Loss,
-                   mode: AMP_TYPE,
-                   amp_config: Config = None):
+def convert_to_amp(model: nn.Module, optimizer: Optimizer, criterion: _Loss, mode: AMP_TYPE, amp_config: Config = None):
     """A helper function to wrap training components with Torch AMP modules
 
     :param model: your model object
     :type model: :class:`torch.nn.Module`
     :param optimizer: your optimizer object
-    :type optimizer: :class:`torch.optim.Optimzer`
+    :type optimizer: :class:`torch.optim.Optimizer`
     :param criterion: your loss function object
     :type criterion: :class:`torch.nn.modules.loss._Loss`
     :param mode: amp mode

colossalai/amp/apex_amp/__init__.py

@@ -3,15 +3,13 @@ import torch.nn as nn
 from torch.optim import Optimizer
 
 
-def convert_to_apex_amp(model: nn.Module,
-                        optimizer: Optimizer,
-                        amp_config):
+def convert_to_apex_amp(model: nn.Module, optimizer: Optimizer, amp_config):
     """A helper function to wrap training components with Apex AMP modules
 
     :param model: your model object
     :type model: :class:`torch.nn.Module`
     :param optimizer: your optimizer object
-    :type optimizer: :class:`torch.optim.Optimzer`
+    :type optimizer: :class:`torch.optim.Optimizer`
     :param amp_config: configuration for nvidia apex
     :type amp_config: :class:`colossalai.context.Config` or dict
 

colossalai/amp/naive_amp/__init__.py

@@ -12,7 +12,7 @@ def convert_to_naive_amp(model: nn.Module, optimizer: Optimizer, amp_config):
     :param model: your model object
     :type model: :class:`torch.nn.Module`
     :param optimizer: your optimizer object
-    :type optimizer: :class:`torch.optim.Optimzer`
+    :type optimizer: :class:`torch.optim.Optimizer`
     :param amp_config: configuration for naive mode amp
     :type amp_config: :class:`colossalai.context.Config` or dict
 

colossalai/amp/torch_amp/__init__.py

@@ -15,7 +15,7 @@ def convert_to_torch_amp(model: nn.Module,
     :param model: your model object
     :type model: :class:`torch.nn.Module`
     :param optimizer: your optimizer object
-    :type optimizer: :class:`torch.optim.Optimzer`
+    :type optimizer: :class:`torch.optim.Optimizer`
     :param criterion: your loss function object
     :type criterion: :class:`torch.nn.modules.loss._Loss`, optional
     :param amp_config: configuration for different amp modes