(self, videos: torch.Tensor)
| 114 | |
| 115 | |
| 116 | def encode(self, videos: torch.Tensor) -> torch.Tensor: |
| 117 | device, dtype = videos[0].device, videos[0].dtype |
| 118 | scale = [self.mean.to(device=device, dtype=dtype), |
| 119 | 1.0 / self.std.to(device=device, dtype=dtype)] |
| 120 | output = [ |
| 121 | self.model.encode(u.unsqueeze(0), scale).float().squeeze(0) |
| 122 | for u in videos |
| 123 | ] |
| 124 | |
| 125 | output = torch.stack(output, dim=0) |
| 126 | return output |
| 127 | |
| 128 | |
| 129 | for i, batch_data in tqdm(enumerate(dataloader), disable=(local_rank != 0)): |