src/model/encoder/encoder.py

from abc import ABC, abstractmethod
from typing import Generic, TypeVar

from torch import nn

from ..types import LVSPM

T = TypeVar("T")


class Encoder(nn.Module, ABC, Generic[T]):
    cfg: T

    def __init__(self, cfg: T) -> None:
        super().__init__()
        self.cfg = cfg

    @abstractmethod
    def forward(
        self,
        context: dict,
        target: dict,
    ) -> LVSPM:
        pass