mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-07-17 08:51:59 +00:00
Fix/format colossalai/engine/paramhooks/(#350)
This commit is contained in:
parent
e83970e3dc
commit
54ee8d1254
@ -1,2 +1,3 @@
|
||||
from ._param_hookmgr import BaseParamHookMgr
|
||||
|
||||
__all__ = ["BaseParamHookMgr"]
|
@ -2,7 +2,9 @@ from typing import Callable, List
|
||||
import torch
|
||||
import functools
|
||||
|
||||
|
||||
class BaseParamHookMgr(object):
|
||||
|
||||
def __init__(self, param_list: List[torch.nn.Parameter]) -> None:
|
||||
r"""
|
||||
register backward hook on every parameters of module
|
||||
|
Loading…
Reference in New Issue
Block a user