github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-26 07:22:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

[NFC] polish amp.naive_amp.grad_scaler code style

Browse Source
This commit is contained in:
zbian 2022-11-09 13:34:19 +08:00 committed by binmakeswell
parent e9635eb493
commit 1559a09fb7
1 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
colossalai/amp/naive_amp/grad_scaler/base_grad_scaler.py
View File

@ -1,12 +1,14 @@
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
import torch
from abc import ABC, abstractmethod
from colossalai.logging import get_dist_logger
from torch import Tensor
from typing import Dict
import torch
from torch import Tensor
from colossalai.logging import get_dist_logger
__all__ = ['BaseGradScaler']