| 122 | example_history = torch.randn(1, 2, 57).cuda().to(torch.float32) |
| 123 | |
| 124 | class WrappedVAEDecode(torch.nn.Module): |
| 125 | |
| 126 | def __init__(self, vae_model, nfuture): |
| 127 | super().__init__() |
| 128 | self.vae = vae_model |
| 129 | self.nfuture = nfuture |
| 130 | |
| 131 | def forward(self, motion_latent: torch.Tensor, |
| 132 | history_motion: torch.Tensor): |
| 133 | return self.vae.decode(motion_latent, |
| 134 | history_motion, |
| 135 | nfuture=self.nfuture) |
| 136 | |
| 137 | class WrappedTRTVAE: |
| 138 |