(x2d, y_local, steps_local)
| 132 | |
| 133 | # Simple Euler sampler for CFM at test time |
| 134 | def euler_sample(x2d, y_local, steps_local): |
| 135 | dt = 1.0 / steps_local |
| 136 | for s in range(steps_local): |
| 137 | t_s = torch.full((gt_3D.size(0), 1, 1, 1), s * dt, device=gt_3D.device, dtype=model_dtype) |
| 138 | v_s = model_3d(x2d, y_local, t_s) |
| 139 | y_local = y_local + dt * v_s |
| 140 | return y_local |
| 141 | |
| 142 | # Start from noise |
| 143 | y = torch.randn(B, F, J, 3, device=gt_3D.device, dtype=model_dtype) |