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.
|
||||
optimizer (Optimizer): Optimizer for updating the parameters.
|
||||
"""
|
||||
|
||||
def __init__(self, model, optimizer):
|
||||
self._model = model
|
||||
self._optimizer = optimizer
|
||||
|
|
|
@ -5,6 +5,7 @@ import torch
|
|||
class BaseOpHook(ABC):
|
||||
"""This class allows users to add customized operations
|
||||
before and after the execution of a PyTorch submodule"""
|
||||
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
|
|
Loading…
Reference in New Issue