(self, imgs, labels)
| 238 | return loss |
| 239 | |
| 240 | def forward(self, imgs, labels): |
| 241 | |
| 242 | # class embed |
| 243 | class_embedding = self.class_emb(labels) |
| 244 | |
| 245 | # patchify and mask (drop) tokens |
| 246 | x = self.patchify(imgs) |
| 247 | gt_latents = x.clone().detach() |
| 248 | orders = self.sample_orders(bsz=x.size(0)) |
| 249 | mask = self.random_masking(x, orders) |
| 250 | |
| 251 | # mae encoder |
| 252 | x = self.forward_mae_encoder(x, mask, class_embedding) |
| 253 | |
| 254 | # mae decoder |
| 255 | z = self.forward_mae_decoder(x, mask) |
| 256 | |
| 257 | # diffloss |
| 258 | loss = self.forward_loss(z=z, target=gt_latents, mask=mask) |
| 259 | |
| 260 | return loss |
| 261 | |
| 262 | def sample_tokens(self, bsz, num_iter=64, cfg=1.0, cfg_schedule="linear", labels=None, temperature=1.0, progress=False): |
| 263 |
nothing calls this directly
no test coverage detected