(
self,
x: torch.Tensor,
action: torch.Tensor,
timestep: torch.Tensor,
context: torch.Tensor,
state: torch.Tensor,
embodiment_id: torch.Tensor,
)
| 557 | self.engine = Engine(eng_path) |
| 558 | |
| 559 | def forward( |
| 560 | self, |
| 561 | x: torch.Tensor, |
| 562 | action: torch.Tensor, |
| 563 | timestep: torch.Tensor, |
| 564 | context: torch.Tensor, |
| 565 | state: torch.Tensor, |
| 566 | embodiment_id: torch.Tensor, |
| 567 | ): |
| 568 | |
| 569 | self.engine.set_runtime_tensor_shape("x", x.shape) |
| 570 | self.engine.set_runtime_tensor_shape("action", action.shape) |
| 571 | self.engine.set_runtime_tensor_shape("context", context.shape) |
| 572 | self.engine.set_runtime_tensor_shape("state", state.shape) |
| 573 | |
| 574 | output = self.engine( |
| 575 | x=x.to(torch.float16), |
| 576 | action=action.to(torch.float16), |
| 577 | timestep=timestep.to(torch.float16), |
| 578 | context=context.to(torch.float16), |
| 579 | state=state.to(torch.float16), |
| 580 | embodiment_id=embodiment_id.to(torch.int32), |
| 581 | ) |
| 582 | if "out.0" in output: # for nvfp4 model export through modelopt |
| 583 | return output["out.0"].to(torch.bfloat16).contiguous(), output["out.1"].to(torch.bfloat16).contiguous() |
| 584 | else: |
| 585 | return output["video_noise_pred"].to(torch.bfloat16).contiguous(), output["action_noise_pred"].to(torch.bfloat16).contiguous() |
| 586 | |
| 587 | |
| 588 | class WanTrtModel14B(torch.nn.Module): |
nothing calls this directly
no test coverage detected