Method __init__

flowmap/model/model_wrapper_overfit.py:25–42

(
        self,
        cfg: ModelWrapperOverfitCfg,
        model: Model,
        batch: Batch,
        flows: Flows,
        tracks: list[Tracks] | None,
        losses: list[Loss],
        visualizers: list[Visualizer],
    )

class ModelWrapperOverfit(LightningModule):
    def __init__(
        self,
        cfg: ModelWrapperOverfitCfg,
        model: Model,
        batch: Batch,
        flows: Flows,
        tracks: list[Tracks] | None,
        losses: list[Loss],
        visualizers: list[Visualizer],
    ) -> None:
        super().__init__()
        self.cfg = cfg
        self.batch = batch
        self.flows = flows
        self.tracks = tracks
        self.model = model
        self.losses = losses
        self.visualizers = visualizers

    def to(self, device: torch.device) -> None:
        self.batch = self.batch.to(device)