mirror of https://github.com/hpcaitech/ColossalAI
fix format (#586)
parent
ffad81e1d1
commit
f2da21a827
|
@ -12,6 +12,7 @@ class BaseGradientHandler(ABC):
|
||||||
model (Module): Model where the gradients accumulate.
|
model (Module): Model where the gradients accumulate.
|
||||||
optimizer (Optimizer): Optimizer for updating the parameters.
|
optimizer (Optimizer): Optimizer for updating the parameters.
|
||||||
"""
|
"""
|
||||||
|
|
||||||
def __init__(self, model, optimizer):
|
def __init__(self, model, optimizer):
|
||||||
self._model = model
|
self._model = model
|
||||||
self._optimizer = optimizer
|
self._optimizer = optimizer
|
||||||
|
|
|
@ -5,6 +5,7 @@ import torch
|
||||||
class BaseOpHook(ABC):
|
class BaseOpHook(ABC):
|
||||||
"""This class allows users to add customized operations
|
"""This class allows users to add customized operations
|
||||||
before and after the execution of a PyTorch submodule"""
|
before and after the execution of a PyTorch submodule"""
|
||||||
|
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
Loading…
Reference in New Issue