(
self, optimizer, clip_grad, log_num_zeros_in_grad, params_have_main_grad
)
| 452 | |
| 453 | class FP32Optimizer(MegatronOptimizer): |
| 454 | def __init__( |
| 455 | self, optimizer, clip_grad, log_num_zeros_in_grad, params_have_main_grad |
| 456 | ): |
| 457 | |
| 458 | super(FP32Optimizer, self).__init__( |
| 459 | optimizer, clip_grad, log_num_zeros_in_grad, params_have_main_grad |
| 460 | ) |
| 461 | |
| 462 | self._scale = torch.cuda.FloatTensor([1.0]) |
| 463 | |
| 464 | def zero_grad(self, set_to_none=True): |
| 465 | """Copied from torch.optim.optimizer""" |