(data_iterator, model, args, timers, only_log_video_latents=False, data_class=None)
| 158 | |
| 159 | |
| 160 | def forward_step_eval(data_iterator, model, args, timers, only_log_video_latents=False, data_class=None): |
| 161 | if mpu.get_model_parallel_rank() == 0: |
| 162 | timers("data loader").start() |
| 163 | batch_video = next(data_iterator) |
| 164 | timers("data loader").stop() |
| 165 | |
| 166 | if len(batch_video["mp4"].shape) == 6: |
| 167 | b, v = batch_video["mp4"].shape[:2] |
| 168 | batch_video["mp4"] = batch_video["mp4"].view(-1, *batch_video["mp4"].shape[2:]) |
| 169 | txt = [] |
| 170 | for i in range(b): |
| 171 | for j in range(v): |
| 172 | txt.append(batch_video["txt"][j][i]) |
| 173 | batch_video["txt"] = txt |
| 174 | |
| 175 | for key in batch_video: |
| 176 | if isinstance(batch_video[key], torch.Tensor): |
| 177 | batch_video[key] = batch_video[key].cuda() |
| 178 | else: |
| 179 | batch_video = {"mp4": None, "fps": None, "num_frames": None, "txt": None} |
| 180 | broad_cast_batch(batch_video) |
| 181 | |
| 182 | |
| 183 | batch_video["global_step"] = args.iteration |
| 184 | loss, loss_dict = model.shared_step(batch_video) |
| 185 | for k in loss_dict: |
| 186 | if loss_dict[k].dtype == torch.bfloat16: |
| 187 | loss_dict[k] = loss_dict[k].to(torch.float32) |
| 188 | return loss, loss_dict |
| 189 | |
| 190 | |
| 191 | def forward_step(data_iterator, model, args, timers, data_class=None): |
nothing calls this directly
no test coverage detected