Proper type hint for init method in class MNISTLitModule(LightningModule)

[cnjackhu]

init method of the class MNISTLitModule(LightningModule):

    def __init__(
        self,
        net: torch.nn.Module,
        optimizer: torch.optim.Optimizer,
        scheduler: torch.optim.lr_scheduler,
        compile: bool,
    ) -> None:

But actually the 'optimizer' and 'scheduler' we pass from a functools.partial, So does the Proper type hint is torch.optim.Optimizer, and torch.optim.lr_scheduler?

The file from model/mnist.yaml

optimizer:
  _target_: torch.optim.Adam
  _partial_: true
  lr: 0.001
  weight_decay: 0.0

scheduler:
  _target_: torch.optim.lr_scheduler.ReduceLROnPlateau
  _partial_: true
  mode: min
  factor: 0.1
  patience: 10