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):
        self.model: nn.Module Optional[nn.Module] = None

    def foo(
            blah: int,
            teacher_lit: pl.LightningModule):
        inputs = torch.rand(3, 2)

mypy output:

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

there are a few more issues in the PL, you are welcome to help :]