Add FSDP support to MNISTLitModule (#604)

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 {