update inference code
Browse files
optim.py
CHANGED
@@ -10,7 +10,7 @@ import torch.nn as nn
|
|
10 |
from torch.distributed.fsdp import FullyShardedDataParallel
|
11 |
from torch.optim.optimizer import Optimizer as OptimizerBase
|
12 |
|
13 |
-
from . import LayerNormBase, BitLinear158
|
14 |
from .config import OptimizerType, SchedulerConfig, SchedulerType, TrainConfig
|
15 |
from .torch_util import get_default_device, is_distributed
|
16 |
|
|
|
10 |
from torch.distributed.fsdp import FullyShardedDataParallel
|
11 |
from torch.optim.optimizer import Optimizer as OptimizerBase
|
12 |
|
13 |
+
from .model import LayerNormBase, BitLinear158
|
14 |
from .config import OptimizerType, SchedulerConfig, SchedulerType, TrainConfig
|
15 |
from .torch_util import get_default_device, is_distributed
|
16 |
|