| 251 | return loss |
| 252 | |
| 253 | def get_loss(self, model_output, target, w): |
| 254 | if self.type == "l2": |
| 255 | # NOTE: Condition frames are excluded from loss calculation |
| 256 | loss = torch.mean((w * (model_output - target) ** 2).reshape(target.shape[0], -1), 1) # * Reshape to [bs] |
| 257 | return loss |
| 258 | elif self.type == "l1": |
| 259 | return torch.mean((w * (model_output - target).abs()).reshape(target.shape[0], -1), 1) |
| 260 | elif self.type == "lpips": |
| 261 | loss = self.lpips(model_output, target).reshape(-1) |
| 262 | return loss |
| 263 | |
| 264 | |
| 265 | def get_dynamics_consistency_loss(self, model_output, target, w, loss_weight=1.): |