diff --git a/reagent/training/multi_stage_trainer.py b/reagent/training/multi_stage_trainer.py
index 5caaec6e..aee79310 100644
--- a/reagent/training/multi_stage_trainer.py
+++ b/reagent/training/multi_stage_trainer.py
@@ -7,7 +7,7 @@ from collections import OrderedDict
 from typing import List, Dict, Tuple
 
 import torch.nn as nn
-from pytorch_lightning.loops.closure import ClosureResult
+from pytorch_lightning.loops.optimization.optimizer_loop import ClosureResult
 from reagent.core.utils import lazy_property
 
 from .reagent_lightning_module import ReAgentLightningModule
diff --git a/setup.cfg b/setup.cfg
index 6ed1e61e..c7ef7804 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -26,7 +26,7 @@ install_requires =
   petastorm>=0.9.0
   parameterized>=0.7.4
   pyspark==3.1.1
-  pytorch-lightning @ git+https://github.com/PyTorchLightning/pytorch-lightning@25af4b137
+  pytorch-lightning @ git+https://github.com/PyTorchLightning/pytorch-lightning@3aba9d16a
   ruamel.yaml>=0.15.99
   scipy>=1.3.1
   tensorboard>=1.14