mirror of https://github.com/hpcaitech/ColossalAI
aibig-modeldata-parallelismdeep-learningdistributed-computingfoundation-modelsheterogeneous-traininghpcinferencelarge-scalemodel-parallelismpipeline-parallelism
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
8 lines
431 B
8 lines
431 B
from torch import Tensor |
|
|
|
class CPUAdamOptimizer: |
|
def __init__(self, lr: float, beta1: float, beta2: float, eps: float, |
|
weight_decay: float, adamw_mode: float) -> None: ... |
|
|
|
def step(self, step: int, lr: float, beta1: float, beta2: float, eps: float, weight_decay: float, bias_correction: bool, |
|
param: Tensor, grad: Tensor, exp_avg: Tensor, exp_avg_sq: Tensor, loss_scale: float) -> None: ...
|
|
|