blog_11.py

from typing import List, Optional, NamedTuple

# use to make model jittable
OptTensor = Optional[Tensor]
ListTensor = List[Tensor]

class TensorBatch(NamedTuple):
    x: Tensor
    edge_index: ListTensor
    edge_attr: OptTensor
    batch: OptTensor
        
class DNAConvNet(LightningModule):
    ...        
        
    def forward(self, batch: TensorBatch):
        ...