| 229 | return loss_value, result, ind_pred.detach(), perplexity.detach() |
| 230 | |
| 231 | def infer_vqstar_completion(self, data, batch_size, mask_ratio, loss_fn='mse'): |
| 232 | # figure out the dimensions: squeeze, permute etc |
| 233 | bev_seq = data['bev_seq'].squeeze(1).permute(0,3,1,2) |
| 234 | bev_seq_next = data['bev_seq_next'] # [bxa, ts-1, H, W, C] |
| 235 | if bev_seq_next.dim()>2: |
| 236 | bev_seq_next = bev_seq_next.permute(0,1,4,2,3) # [bxa, ts-1, C, H, W] |
| 237 | # print(data['bev_seq_teacher'].size()) |
| 238 | bev_teacher = data['bev_seq_teacher'].squeeze(1).permute(0,3,1,2) |
| 239 | # print("teacher size", bev_teacher.size()) |
| 240 | num_agent_tensor = data['num_agent'] |
| 241 | trans_matrices = data['trans_matrices'] |
| 242 | |
| 243 | with torch.cuda.amp.autocast(enabled=False): |
| 244 | loss, result, ind_pred, perplexity, encodings = self.model.inference(bev_seq, |
| 245 | bev_seq_next, |
| 246 | bev_teacher, |
| 247 | trans_matrices, |
| 248 | num_agent_tensor, |
| 249 | batch_size, |
| 250 | mask_ratio=mask_ratio) |
| 251 | |
| 252 | loss_value = loss.item() |
| 253 | return loss_value, result, ind_pred.detach(), perplexity.detach(), encodings.detach() |