Lightning as the default runner in e2e workflow

Summary: As titled.

Reviewed By: zhanghang1989

Differential Revision: D29075952

fbshipit-source-id: 6ef3dc35cd436c1fffb031ea59f20ca23afc5368

d2go/runner/lightning_task.py

@@ -53,7 +53,7 @@ def _is_lightning_checkpoint(checkpoint: Dict[str, Any]) -> bool:
 
 def _is_d2go_checkpoint(checkpoint: Dict[str, Any]) -> bool:
     """Returns true if we believe this to be a D2Go checkpoint."""
-    d2_go_keys = [_OLD_STATE_DICT_KEY, "optimizer", "scheduler", "iteration"]
+    d2_go_keys = [_OLD_STATE_DICT_KEY, "iteration"]
     for key in d2_go_keys:
         if key not in checkpoint:
             return False
@@ -85,6 +85,8 @@ def _convert_to_lightning(d2_checkpoint: Dict[str, Any]) -> None:
     for old, new in zip(
         ["optimizer", "scheduler"], ["optimizer_states", "lr_schedulers"]
     ):
+        if old not in d2_checkpoint:
+            continue
         d2_checkpoint[new] = [d2_checkpoint[old]]
         del d2_checkpoint[old]