fix optimizer setting in pytorch lightning

Summary: fix optimizer setting in pytorch lightning

Reviewed By: wat3rBro

Differential Revision: D30988441

fbshipit-source-id: fcd2f4c77a87a790d7e99b0e3c833c291fd66e77

d2go/runner/lightning_task.py

@@ -23,6 +23,7 @@ from d2go.modeling.quantization import (
     default_prepare_for_quant,
     default_prepare_for_quant_convert,
 )
+from d2go.optimizer import build_optimizer_mapper
 from d2go.runner.callbacks.quantization import maybe_prepare_for_quantization, PREPARED
 from d2go.runner.default_runner import (
     Detectron2GoRunner,
@@ -225,7 +226,7 @@ class DefaultTask(pl.LightningModule):
     def configure_optimizers(
         self,
     ) -> Tuple[List[torch.optim.Optimizer], List]:
-        optim = d2_build_optimizer(self.cfg, self.model)
+        optim = build_optimizer_mapper(self.cfg, self.model)
         lr_scheduler = d2_build_lr_scheduler(self.cfg, optim)
 
         return [optim], [{"scheduler": lr_scheduler, "interval": "step"}]