Mypy issues with lightning module

I’m using mypy with lightning. However, the types associated with attributes assigned to a lightning module seem not to correspond with the types of the objects I’m assigning. How to handle this?

A minimum reproduceable example:

import pytorch_lightning as pl
from typing import Optional, cast
import torch
from torch import nn


class MyLightningModule(pl.LightningModule):
    def __init__(self):
        super.__init__()
        self.model: nn.Module
        self.foo: Optional[nn.Module] = None

    def foo(
            self,
            blah: int,
            teacher_lit: pl.LightningModule):
        super().__init__()
        inputs = torch.rand(3, 2)
        self.model(inputs)
        cast(nn.Module, self.foo)(inputs)
        reveal_type(self.foo)
        self.foo(inputs)
        self.bar(inputs)

mypy output:

(asr) asr ~/git/hugh-pub (master|…5△1)$ mypy ~/git/hugh-pub/prot/test_lightning_inherit.py 
prot/test_lightning_inherit.py:21: note: Revealed type is 'def (blah: builtins.int, teacher_lit: pytorch_lightning.core.lightning.LightningModule) -> Any'
prot/test_lightning_inherit.py:22: error: Too few arguments for "foo" of "MyLightningModule"
prot/test_lightning_inherit.py:22: error: Argument 1 to "foo" of "MyLightningModule" has incompatible type "Tensor"; expected "int"
prot/test_lightning_inherit.py:23: error: "Tensor" not callable

there are a few more issues in the PL, you are welcome to help :]
https://github.com/PyTorchLightning/pytorch-lightning/issues/5023