fix gradient clipping
Summary: Pull Request resolved:
@@ -112,14 +112,14 @@ class Trainer(DefaultTrainer):
optimizer_type = cfg.SOLVER.OPTIMIZER
if optimizer_type == "SGD":
- return maybe_add_gradient_clipping(torch.optim.SGD)(
+ return maybe_add_gradient_clipping(cfg, torch.optim.SGD)(
params,
cfg.SOLVER.BASE_LR,
momentum=cfg.SOLVER.MOMENTUM,
nesterov=cfg.SOLVER.NESTEROV,
)
elif optimizer_type == "ADAM":
- return maybe_add_gradient_clipping(torch.optim.Adam)(params, cfg.SOLVER.BASE_LR)
+ return maybe_add_gradient_clipping(cfg, torch.optim.Adam)(params, cfg.SOLVER.BASE_LR)
else:
raise NotImplementedError(f"no optimizer type {optimizer_type}")
