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