From bddbc24b82ab6ccfa6243e815a49dc5bfe8d4144 Mon Sep 17 00:00:00 2001
From: Alex Morehead <alex.morehead@gmail.com>
Date: Mon, 25 Sep 2023 15:23:58 -0500
Subject: [PATCH] Add FSDP support to `MNISTLitModule` (#604)

---
 src/models/mnist_module.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/models/mnist_module.py b/src/models/mnist_module.py
index b10450f19..5d303ac2f 100644
--- a/src/models/mnist_module.py
+++ b/src/models/mnist_module.py
@@ -198,7 +198,7 @@ def configure_optimizers(self) -> Dict[str, Any]:
 
         :return: A dict containing the configured optimizers and learning-rate schedulers to be used for training.
         """
-        optimizer = self.hparams.optimizer(params=self.parameters())
+        optimizer = self.hparams.optimizer(params=self.trainer.model.parameters())
         if self.hparams.scheduler is not None:
             scheduler = self.hparams.scheduler(optimizer=optimizer)
             return {