(
self,
cfg: ModelWrapperPretrainCfg,
cfg_cropping: CroppingCfg,
cfg_flow: FlowPredictorCfg,
model: Model,
losses: list[Loss],
visualizers: list[Visualizer],
)
| 27 | |
| 28 | class ModelWrapperPretrain(LightningModule): |
| 29 | def __init__( |
| 30 | self, |
| 31 | cfg: ModelWrapperPretrainCfg, |
| 32 | cfg_cropping: CroppingCfg, |
| 33 | cfg_flow: FlowPredictorCfg, |
| 34 | model: Model, |
| 35 | losses: list[Loss], |
| 36 | visualizers: list[Visualizer], |
| 37 | ) -> None: |
| 38 | super().__init__() |
| 39 | self.cfg = cfg |
| 40 | self.cfg_cropping = cfg_cropping |
| 41 | self.flow_predictor = get_flow_predictor(cfg_flow) |
| 42 | self.model = model |
| 43 | self.losses = losses |
| 44 | self.visualizers = visualizers |
| 45 | |
| 46 | @torch.no_grad() |
| 47 | def preprocess_batch(self, batch_dict: dict) -> tuple[Batch, Flows]: |
nothing calls this directly
no test coverage detected