We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `colossalai/nn/optimizer/cpu_adam.py`
Content:
```
1 import math
2 from typing import Optional
3 
4 import torch
5 
6 from colossalai.kernel.op_builder import CPUAdamBuilder
7 
8 from .nvme_optimizer import NVMeOptimizer
9 
10 
11 class CPUAdam(NVMeOptimizer):
12     """Implements Adam algorithm.
13 
14     Supports parameters updating on both GPU and CPU, depending on the device of parameters.
15     But the parameters and gradients should on the same device:
16       * Parameters on CPU and gradients on CPU is allowed.
17       * Parameters on GPU and gradients on GPU is allowed.
18       * Parameters on GPU and gradients on CPU is **not** allowed.
19 
20     `CPUAdam` requires CUDA extensions which can be built during installation or runtime.
21 
22     This version of CPU Adam accelerates parameters updating on CPU with SIMD.
23     Support of AVX2 or AVX512 is required.
24 
25     The GPU part is implemented in an naive way.
26 
27     CPU Adam also supports the hybrid precision calculation, eg. fp32 parameters and fp16 gradients.
28 
29     :class:`colossalai.nn.optimizer.CPUAdam` may be used as a drop-in replacement for ``torch.optim.AdamW``,
30     or ``torch.optim.Adam`` with ``adamw_mode=False``
31 
32     Adam was been proposed in `Adam: A Method for Stochastic Optimization`_.
33 
34     Arguments:
35         model_params (iterable): iterable of parameters of dicts defining
36             parameter groups.
37         lr (float, optional): learning rate. (default: 1e-3)
38         betas (Tuple[float, float], optional): coefficients used for computing
39             running averages of gradient and its square. (default: (0.9, 0.999))
40         eps (float, optional): term added to the denominator to improve
41             numerical stability. (default: 1e-8)
42         weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
43         amsgrad (boolean, optional): whether to use the AMSGrad variant of this
44             algorithm from the paper `On the Convergence of Adam and Beyond`_
45             (default: False) NOT SUPPORTED yet in CPUAdam!
46         adamw_mode (boolean, optional): Apply L2 regularization or weight decay
47             True for decoupled weight decay(also known as AdamW) (default: True)
48         simd_log (boolean, optional): whether to show if you are using SIMD to
49             accelerate. (default: False)
50         nvme_offload_fraction (float, optional): Fraction of optimizer states to be offloaded to NVMe. Defaults to 0.0.
51         nvme_offload_dir (Optional[str], optional): Directory to save NVMe offload files.
52             If it's ``None``, a random temporary directory will be used. Defaults to None.
53 
54     .. _Adam\: A Method for Stochastic Optimization:
55         https://arxiv.org/abs/1412.6980
56     .. _On the Convergence of Adam and Beyond:
57         https://openreview.net/forum?id=ryQu7f-RZ
58     """
59 
60     # Number of fp32 shards for per parameter
61     # Param weight, grad, momentum and variance
62     num_fp32_shards_per_param = 4
63 
64     def __init__(
65         self,
66         model_params,
67         lr=1e-3,
68         bias_correction=True,
69         betas=(0.9, 0.999),
70         eps=1e-8,
71         weight_decay=0,
72         adamw_mode=True,
73         nvme_offload_fraction: float = 0.0,
74         nvme_offload_dir: Optional[str] = None,
75     ):
76         default_args = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, bias_correction=bias_correction)
77         super(CPUAdam, self).__init__(model_params, default_args, nvme_offload_fraction, nvme_offload_dir)
78         self.adamw_mode = adamw_mode
79         cpu_adam = CPUAdamBuilder().load()
80         # if you find yourself stuck here, make sure that you install colossalai with CUDA_EXT=1 specification
81         self.cpu_adam_op = cpu_adam.CPUAdamOptimizer(lr, betas[0], betas[1], eps, weight_decay, adamw_mode)
82 
83     def torch_adam_update(
84         self,
85         data,
86         grad,
87         exp_avg,
88         exp_avg_sq,
89         lr,
90         beta1,
91         beta2,
92         eps,
93         weight_decay,
94         bias_correction1,
95         bias_correction2,
96         use_adamw=False,
97     ):
98         grad = grad.to(data.dtype)
99 
100         if weight_decay != 0:
101             if use_adamw:
102                 data.mul_(1 - lr * weight_decay)
103             else:
104                 grad = grad.add(data, alpha=weight_decay)
105 
106         # Decay the first and second moment running average coefficient
107         exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
108         exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
109 
110         # TODO(jiaruifang) dose not support amsgrad
111         denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
112 
113         step_size = lr / bias_correction1
114 
115         data.addcdiv_(exp_avg, denom, value=-step_size)
116 
117     @torch.no_grad()
118     def step(self, closure=None, div_scale: float = -1):
119         loss = None
120         if closure is not None:
121             with torch.enable_grad():
122                 loss = closure()
123 
124         self._pre_step("exp_avg", "exp_avg_sq")
125         for _, group in enumerate(self.param_groups):
126             for _, p in enumerate(group["params"]):
127                 if p.grad is None:
128                     continue
129 
130                 state = self.state[p]
131 
132                 target_device = p.device
133                 if len(state) == 0:
134                     state["step"] = 0
135                     # gradient momentums
136                     state["exp_avg"] = torch.zeros_like(p, device=target_device)
137                     # gradient variances
138                     state["exp_avg_sq"] = torch.zeros_like(p, device=target_device)
139                     self._post_state_init(p)
140 
141                 state["step"] += 1
142                 beta1, beta2 = group["betas"]
143 
144                 if target_device.type == "cpu":
145                     assert p.data.numel() == p.grad.data.numel(), "parameter and gradient should have the same size"
146                     assert state["exp_avg"].device.type == "cpu", "exp_avg should stay on cpu"
147                     assert state["exp_avg_sq"].device.type == "cpu", "exp_avg should stay on cpu"
148                     self._pre_update(p, "exp_avg", "exp_avg_sq")
149                     if p.grad.dtype is torch.bfloat16:
150                         # cpu adam kernel does not support bf16 now
151                         bias_correction1 = 1 - beta1 ** state["step"]
152                         bias_correction2 = 1 - beta2 ** state["step"]
153                         self.torch_adam_update(
154                             p.data,
155                             p.grad.data,
156                             state["exp_avg"],
157                             state["exp_avg_sq"],
158                             group["lr"],
159                             beta1,
160                             beta2,
161                             group["eps"],
162                             group["weight_decay"],
163                             bias_correction1,
164                             bias_correction2,
165                             self.adamw_mode,
166                         )
167                     else:
168                         self.cpu_adam_op.step(
169                             state["step"],
170                             group["lr"],
171                             beta1,
172                             beta2,
173                             group["eps"],
174                             group["weight_decay"],
175                             group["bias_correction"],
176                             p.data,
177                             p.grad.data,
178                             state["exp_avg"],
179                             state["exp_avg_sq"],
180                             div_scale,
181                         )
182                     self._post_update(p, "exp_avg", "exp_avg_sq")
183                 elif target_device.type == "cuda":
184                     assert div_scale == -1, "div_scale should remain default"
185                     assert state["exp_avg"].device.type == "cuda", "exp_avg should stay on cuda"
186                     assert state["exp_avg_sq"].device.type == "cuda", "exp_avg should stay on cuda"
187 
188                     bias_correction1 = 1 - beta1 ** state["step"]
189                     bias_correction2 = 1 - beta2 ** state["step"]
190 
191                     # adam on cuda
192                     self.torch_adam_update(
193                         p.data,
194                         p.grad.data,
195                         state["exp_avg"],
196                         state["exp_avg_sq"],
197                         group["lr"],
198                         beta1,
199                         beta2,
200                         group["eps"],
201                         group["weight_decay"],
202                         bias_correction1,
203                         bias_correction2,
204                         self.adamw_mode,
205                     )
206                 else:
207                     raise RuntimeError
208         self._post_step()
209         return loss
210
```
--- 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/colossalai/nn/optimizer/cpu_adam.py b/colossalai/nn/optimizer/cpu_adam.py
--- a/colossalai/nn/optimizer/cpu_adam.py
+++ b/colossalai/nn/optimizer/cpu_adam.py
@@ -9,7 +9,8 @@
 
 
 class CPUAdam(NVMeOptimizer):
-    """Implements Adam algorithm.
+    """
+    Implements Adam algorithm.
 
     Supports parameters updating on both GPU and CPU, depending on the device of parameters.
     But the parameters and gradients should on the same device:

{"golden_diff": "diff --git a/colossalai/nn/optimizer/cpu_adam.py b/colossalai/nn/optimizer/cpu_adam.py\n--- a/colossalai/nn/optimizer/cpu_adam.py\n+++ b/colossalai/nn/optimizer/cpu_adam.py\n@@ -9,7 +9,8 @@\n \n \n class CPUAdam(NVMeOptimizer):\n-    \"\"\"Implements Adam algorithm.\n+    \"\"\"\n+    Implements Adam algorithm.\n \n     Supports parameters updating on both GPU and CPU, depending on the device of parameters.\n     But the parameters and gradients should on the same device:\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "import math\nfrom typing import Optional\n\nimport torch\n\nfrom colossalai.kernel.op_builder import CPUAdamBuilder\n\nfrom .nvme_optimizer import NVMeOptimizer\n\n\nclass CPUAdam(NVMeOptimizer):\n    \"\"\"Implements Adam algorithm.\n\n    Supports parameters updating on both GPU and CPU, depending on the device of parameters.\n    But the parameters and gradients should on the same device:\n      * Parameters on CPU and gradients on CPU is allowed.\n      * Parameters on GPU and gradients on GPU is allowed.\n      * Parameters on GPU and gradients on CPU is **not** allowed.\n\n    `CPUAdam` requires CUDA extensions which can be built during installation or runtime.\n\n    This version of CPU Adam accelerates parameters updating on CPU with SIMD.\n    Support of AVX2 or AVX512 is required.\n\n    The GPU part is implemented in an naive way.\n\n    CPU Adam also supports the hybrid precision calculation, eg. fp32 parameters and fp16 gradients.\n\n    :class:`colossalai.nn.optimizer.CPUAdam` may be used as a drop-in replacement for ``torch.optim.AdamW``,\n    or ``torch.optim.Adam`` with ``adamw_mode=False``\n\n    Adam was been proposed in `Adam: A Method for Stochastic Optimization`_.\n\n    Arguments:\n        model_params (iterable): iterable of parameters of dicts defining\n            parameter groups.\n        lr (float, optional): learning rate. (default: 1e-3)\n        betas (Tuple[float, float], optional): coefficients used for computing\n            running averages of gradient and its square. (default: (0.9, 0.999))\n        eps (float, optional): term added to the denominator to improve\n            numerical stability. (default: 1e-8)\n        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)\n        amsgrad (boolean, optional): whether to use the AMSGrad variant of this\n            algorithm from the paper `On the Convergence of Adam and Beyond`_\n            (default: False) NOT SUPPORTED yet in CPUAdam!\n        adamw_mode (boolean, optional): Apply L2 regularization or weight decay\n            True for decoupled weight decay(also known as AdamW) (default: True)\n        simd_log (boolean, optional): whether to show if you are using SIMD to\n            accelerate. (default: False)\n        nvme_offload_fraction (float, optional): Fraction of optimizer states to be offloaded to NVMe. Defaults to 0.0.\n        nvme_offload_dir (Optional[str], optional): Directory to save NVMe offload files.\n            If it's ``None``, a random temporary directory will be used. Defaults to None.\n\n    .. _Adam\\: A Method for Stochastic Optimization:\n        https://arxiv.org/abs/1412.6980\n    .. _On the Convergence of Adam and Beyond:\n        https://openreview.net/forum?id=ryQu7f-RZ\n    \"\"\"\n\n    # Number of fp32 shards for per parameter\n    # Param weight, grad, momentum and variance\n    num_fp32_shards_per_param = 4\n\n    def __init__(\n        self,\n        model_params,\n        lr=1e-3,\n        bias_correction=True,\n        betas=(0.9, 0.999),\n        eps=1e-8,\n        weight_decay=0,\n        adamw_mode=True,\n        nvme_offload_fraction: float = 0.0,\n        nvme_offload_dir: Optional[str] = None,\n    ):\n        default_args = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, bias_correction=bias_correction)\n        super(CPUAdam, self).__init__(model_params, default_args, nvme_offload_fraction, nvme_offload_dir)\n        self.adamw_mode = adamw_mode\n        cpu_adam = CPUAdamBuilder().load()\n        # if you find yourself stuck here, make sure that you install colossalai with CUDA_EXT=1 specification\n        self.cpu_adam_op = cpu_adam.CPUAdamOptimizer(lr, betas[0], betas[1], eps, weight_decay, adamw_mode)\n\n    def torch_adam_update(\n        self,\n        data,\n        grad,\n        exp_avg,\n        exp_avg_sq,\n        lr,\n        beta1,\n        beta2,\n        eps,\n        weight_decay,\n        bias_correction1,\n        bias_correction2,\n        use_adamw=False,\n    ):\n        grad = grad.to(data.dtype)\n\n        if weight_decay != 0:\n            if use_adamw:\n                data.mul_(1 - lr * weight_decay)\n            else:\n                grad = grad.add(data, alpha=weight_decay)\n\n        # Decay the first and second moment running average coefficient\n        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)\n        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)\n\n        # TODO(jiaruifang) dose not support amsgrad\n        denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)\n\n        step_size = lr / bias_correction1\n\n        data.addcdiv_(exp_avg, denom, value=-step_size)\n\n    @torch.no_grad()\n    def step(self, closure=None, div_scale: float = -1):\n        loss = None\n        if closure is not None:\n            with torch.enable_grad():\n                loss = closure()\n\n        self._pre_step(\"exp_avg\", \"exp_avg_sq\")\n        for _, group in enumerate(self.param_groups):\n            for _, p in enumerate(group[\"params\"]):\n                if p.grad is None:\n                    continue\n\n                state = self.state[p]\n\n                target_device = p.device\n                if len(state) == 0:\n                    state[\"step\"] = 0\n                    # gradient momentums\n                    state[\"exp_avg\"] = torch.zeros_like(p, device=target_device)\n                    # gradient variances\n                    state[\"exp_avg_sq\"] = torch.zeros_like(p, device=target_device)\n                    self._post_state_init(p)\n\n                state[\"step\"] += 1\n                beta1, beta2 = group[\"betas\"]\n\n                if target_device.type == \"cpu\":\n                    assert p.data.numel() == p.grad.data.numel(), \"parameter and gradient should have the same size\"\n                    assert state[\"exp_avg\"].device.type == \"cpu\", \"exp_avg should stay on cpu\"\n                    assert state[\"exp_avg_sq\"].device.type == \"cpu\", \"exp_avg should stay on cpu\"\n                    self._pre_update(p, \"exp_avg\", \"exp_avg_sq\")\n                    if p.grad.dtype is torch.bfloat16:\n                        # cpu adam kernel does not support bf16 now\n                        bias_correction1 = 1 - beta1 ** state[\"step\"]\n                        bias_correction2 = 1 - beta2 ** state[\"step\"]\n                        self.torch_adam_update(\n                            p.data,\n                            p.grad.data,\n                            state[\"exp_avg\"],\n                            state[\"exp_avg_sq\"],\n                            group[\"lr\"],\n                            beta1,\n                            beta2,\n                            group[\"eps\"],\n                            group[\"weight_decay\"],\n                            bias_correction1,\n                            bias_correction2,\n                            self.adamw_mode,\n                        )\n                    else:\n                        self.cpu_adam_op.step(\n                            state[\"step\"],\n                            group[\"lr\"],\n                            beta1,\n                            beta2,\n                            group[\"eps\"],\n                            group[\"weight_decay\"],\n                            group[\"bias_correction\"],\n                            p.data,\n                            p.grad.data,\n                            state[\"exp_avg\"],\n                            state[\"exp_avg_sq\"],\n                            div_scale,\n                        )\n                    self._post_update(p, \"exp_avg\", \"exp_avg_sq\")\n                elif target_device.type == \"cuda\":\n                    assert div_scale == -1, \"div_scale should remain default\"\n                    assert state[\"exp_avg\"].device.type == \"cuda\", \"exp_avg should stay on cuda\"\n                    assert state[\"exp_avg_sq\"].device.type == \"cuda\", \"exp_avg should stay on cuda\"\n\n                    bias_correction1 = 1 - beta1 ** state[\"step\"]\n                    bias_correction2 = 1 - beta2 ** state[\"step\"]\n\n                    # adam on cuda\n                    self.torch_adam_update(\n                        p.data,\n                        p.grad.data,\n                        state[\"exp_avg\"],\n                        state[\"exp_avg_sq\"],\n                        group[\"lr\"],\n                        beta1,\n                        beta2,\n                        group[\"eps\"],\n                        group[\"weight_decay\"],\n                        bias_correction1,\n                        bias_correction2,\n                        self.adamw_mode,\n                    )\n                else:\n                    raise RuntimeError\n        self._post_step()\n        return loss\n", "path": "colossalai/nn/optimizer/cpu_adam.py"}], "after_files": [{"content": "import math\nfrom typing import Optional\n\nimport torch\n\nfrom colossalai.kernel.op_builder import CPUAdamBuilder\n\nfrom .nvme_optimizer import NVMeOptimizer\n\n\nclass CPUAdam(NVMeOptimizer):\n    \"\"\"\n    Implements Adam algorithm.\n\n    Supports parameters updating on both GPU and CPU, depending on the device of parameters.\n    But the parameters and gradients should on the same device:\n      * Parameters on CPU and gradients on CPU is allowed.\n      * Parameters on GPU and gradients on GPU is allowed.\n      * Parameters on GPU and gradients on CPU is **not** allowed.\n\n    `CPUAdam` requires CUDA extensions which can be built during installation or runtime.\n\n    This version of CPU Adam accelerates parameters updating on CPU with SIMD.\n    Support of AVX2 or AVX512 is required.\n\n    The GPU part is implemented in an naive way.\n\n    CPU Adam also supports the hybrid precision calculation, eg. fp32 parameters and fp16 gradients.\n\n    :class:`colossalai.nn.optimizer.CPUAdam` may be used as a drop-in replacement for ``torch.optim.AdamW``,\n    or ``torch.optim.Adam`` with ``adamw_mode=False``\n\n    Adam was been proposed in `Adam: A Method for Stochastic Optimization`_.\n\n    Arguments:\n        model_params (iterable): iterable of parameters of dicts defining\n            parameter groups.\n        lr (float, optional): learning rate. (default: 1e-3)\n        betas (Tuple[float, float], optional): coefficients used for computing\n            running averages of gradient and its square. (default: (0.9, 0.999))\n        eps (float, optional): term added to the denominator to improve\n            numerical stability. (default: 1e-8)\n        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)\n        amsgrad (boolean, optional): whether to use the AMSGrad variant of this\n            algorithm from the paper `On the Convergence of Adam and Beyond`_\n            (default: False) NOT SUPPORTED yet in CPUAdam!\n        adamw_mode (boolean, optional): Apply L2 regularization or weight decay\n            True for decoupled weight decay(also known as AdamW) (default: True)\n        simd_log (boolean, optional): whether to show if you are using SIMD to\n            accelerate. (default: False)\n        nvme_offload_fraction (float, optional): Fraction of optimizer states to be offloaded to NVMe. Defaults to 0.0.\n        nvme_offload_dir (Optional[str], optional): Directory to save NVMe offload files.\n            If it's ``None``, a random temporary directory will be used. Defaults to None.\n\n    .. _Adam\\: A Method for Stochastic Optimization:\n        https://arxiv.org/abs/1412.6980\n    .. _On the Convergence of Adam and Beyond:\n        https://openreview.net/forum?id=ryQu7f-RZ\n    \"\"\"\n\n    # Number of fp32 shards for per parameter\n    # Param weight, grad, momentum and variance\n    num_fp32_shards_per_param = 4\n\n    def __init__(\n        self,\n        model_params,\n        lr=1e-3,\n        bias_correction=True,\n        betas=(0.9, 0.999),\n        eps=1e-8,\n        weight_decay=0,\n        adamw_mode=True,\n        nvme_offload_fraction: float = 0.0,\n        nvme_offload_dir: Optional[str] = None,\n    ):\n        default_args = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, bias_correction=bias_correction)\n        super(CPUAdam, self).__init__(model_params, default_args, nvme_offload_fraction, nvme_offload_dir)\n        self.adamw_mode = adamw_mode\n        cpu_adam = CPUAdamBuilder().load()\n        # if you find yourself stuck here, make sure that you install colossalai with CUDA_EXT=1 specification\n        self.cpu_adam_op = cpu_adam.CPUAdamOptimizer(lr, betas[0], betas[1], eps, weight_decay, adamw_mode)\n\n    def torch_adam_update(\n        self,\n        data,\n        grad,\n        exp_avg,\n        exp_avg_sq,\n        lr,\n        beta1,\n        beta2,\n        eps,\n        weight_decay,\n        bias_correction1,\n        bias_correction2,\n        use_adamw=False,\n    ):\n        grad = grad.to(data.dtype)\n\n        if weight_decay != 0:\n            if use_adamw:\n                data.mul_(1 - lr * weight_decay)\n            else:\n                grad = grad.add(data, alpha=weight_decay)\n\n        # Decay the first and second moment running average coefficient\n        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)\n        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)\n\n        # TODO(jiaruifang) dose not support amsgrad\n        denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)\n\n        step_size = lr / bias_correction1\n\n        data.addcdiv_(exp_avg, denom, value=-step_size)\n\n    @torch.no_grad()\n    def step(self, closure=None, div_scale: float = -1):\n        loss = None\n        if closure is not None:\n            with torch.enable_grad():\n                loss = closure()\n\n        self._pre_step(\"exp_avg\", \"exp_avg_sq\")\n        for _, group in enumerate(self.param_groups):\n            for _, p in enumerate(group[\"params\"]):\n                if p.grad is None:\n                    continue\n\n                state = self.state[p]\n\n                target_device = p.device\n                if len(state) == 0:\n                    state[\"step\"] = 0\n                    # gradient momentums\n                    state[\"exp_avg\"] = torch.zeros_like(p, device=target_device)\n                    # gradient variances\n                    state[\"exp_avg_sq\"] = torch.zeros_like(p, device=target_device)\n                    self._post_state_init(p)\n\n                state[\"step\"] += 1\n                beta1, beta2 = group[\"betas\"]\n\n                if target_device.type == \"cpu\":\n                    assert p.data.numel() == p.grad.data.numel(), \"parameter and gradient should have the same size\"\n                    assert state[\"exp_avg\"].device.type == \"cpu\", \"exp_avg should stay on cpu\"\n                    assert state[\"exp_avg_sq\"].device.type == \"cpu\", \"exp_avg should stay on cpu\"\n                    self._pre_update(p, \"exp_avg\", \"exp_avg_sq\")\n                    if p.grad.dtype is torch.bfloat16:\n                        # cpu adam kernel does not support bf16 now\n                        bias_correction1 = 1 - beta1 ** state[\"step\"]\n                        bias_correction2 = 1 - beta2 ** state[\"step\"]\n                        self.torch_adam_update(\n                            p.data,\n                            p.grad.data,\n                            state[\"exp_avg\"],\n                            state[\"exp_avg_sq\"],\n                            group[\"lr\"],\n                            beta1,\n                            beta2,\n                            group[\"eps\"],\n                            group[\"weight_decay\"],\n                            bias_correction1,\n                            bias_correction2,\n                            self.adamw_mode,\n                        )\n                    else:\n                        self.cpu_adam_op.step(\n                            state[\"step\"],\n                            group[\"lr\"],\n                            beta1,\n                            beta2,\n                            group[\"eps\"],\n                            group[\"weight_decay\"],\n                            group[\"bias_correction\"],\n                            p.data,\n                            p.grad.data,\n                            state[\"exp_avg\"],\n                            state[\"exp_avg_sq\"],\n                            div_scale,\n                        )\n                    self._post_update(p, \"exp_avg\", \"exp_avg_sq\")\n                elif target_device.type == \"cuda\":\n                    assert div_scale == -1, \"div_scale should remain default\"\n                    assert state[\"exp_avg\"].device.type == \"cuda\", \"exp_avg should stay on cuda\"\n                    assert state[\"exp_avg_sq\"].device.type == \"cuda\", \"exp_avg should stay on cuda\"\n\n                    bias_correction1 = 1 - beta1 ** state[\"step\"]\n                    bias_correction2 = 1 - beta2 ** state[\"step\"]\n\n                    # adam on cuda\n                    self.torch_adam_update(\n                        p.data,\n                        p.grad.data,\n                        state[\"exp_avg\"],\n                        state[\"exp_avg_sq\"],\n                        group[\"lr\"],\n                        beta1,\n                        beta2,\n                        group[\"eps\"],\n                        group[\"weight_decay\"],\n                        bias_correction1,\n                        bias_correction2,\n                        self.adamw_mode,\n                    )\n                else:\n                    raise RuntimeError\n        self._post_step()\n        return loss\n", "path": "colossalai/nn/optimizer/cpu_adam.py"}]}