(self, motion, motion_mask)
| 67 | self.out = nn.Linear(self.latent_dim, 512) |
| 68 | |
| 69 | def forward(self, motion, motion_mask): |
| 70 | x, mask = motion, motion_mask |
| 71 | B, T = x.shape[:2] |
| 72 | |
| 73 | x = x.reshape(B, T, 2, -1)[..., :-4].reshape(B, T, -1) |
| 74 | |
| 75 | x_emb = self.embed_motion(x) |
| 76 | |
| 77 | idx = torch.zeros(B, dtype=torch.long, device=x.device) |
| 78 | emb = torch.cat([self.query_token[idx][:, None], x_emb], dim=1) |
| 79 | |
| 80 | seq_mask = (mask > 0.5) |
| 81 | token_mask = torch.ones((B, 1), dtype=bool, device=x.device) |
| 82 | valid_mask = torch.cat([token_mask, seq_mask], dim=1) |
| 83 | |
| 84 | h = self.sequence_pos_encoder(emb) |
| 85 | |
| 86 | h = h.permute(1, 0, 2) |
| 87 | h = self.transformer(h, src_key_padding_mask=~valid_mask).permute( |
| 88 | 1, 0, 2) |
| 89 | h = self.out_ln(h) |
| 90 | motion_emb = self.out(h[:, 0]) |
| 91 | |
| 92 | return motion_emb |
| 93 | |
| 94 | |
| 95 | @SUBMODULES.register_module() |
nothing calls this directly
no outgoing calls
no test coverage detected