(
self,
batch: Batch,
flows: Flows,
global_step: int,
)
| 91 | |
| 92 | @torch.no_grad() |
| 93 | def export( |
| 94 | self, |
| 95 | batch: Batch, |
| 96 | flows: Flows, |
| 97 | global_step: int, |
| 98 | ) -> ModelExports: |
| 99 | # For now, only implement exporting with a batch size of 1. |
| 100 | b, _, _, _, _ = batch.videos.shape |
| 101 | assert b == 1 |
| 102 | |
| 103 | output = self.forward(batch, flows, global_step) |
| 104 | |
| 105 | return ModelExports( |
| 106 | output.extrinsics, |
| 107 | output.intrinsics, |
| 108 | batch.videos, |
| 109 | output.depths, |
| 110 | ) |
nothing calls this directly
no test coverage detected