(
self,
batch: Batch,
flows: Flows,
tracks: list[Tracks] | None,
model_output: ModelOutput,
global_step: int,
)
| 29 | self.cfg = cfg |
| 30 | |
| 31 | def forward( |
| 32 | self, |
| 33 | batch: Batch, |
| 34 | flows: Flows, |
| 35 | tracks: list[Tracks] | None, |
| 36 | model_output: ModelOutput, |
| 37 | global_step: int, |
| 38 | ) -> Float[Tensor, ""]: |
| 39 | # Before the loss is enabled, don't compute the loss. |
| 40 | if global_step < self.cfg.enable_after: |
| 41 | return torch.tensor(0, dtype=torch.float32, device=batch.videos.device) |
| 42 | |
| 43 | # Multiply the computed loss value by the weight. |
| 44 | loss = self.compute_unweighted_loss( |
| 45 | batch, flows, tracks, model_output, global_step |
| 46 | ) |
| 47 | return self.cfg.weight * loss |
| 48 | |
| 49 | @abstractmethod |
| 50 | def compute_unweighted_loss( |
no test coverage detected