(self, probs)
| 251 | |
| 252 | @torch.no_grad() |
| 253 | def calculate_mask(self, probs): |
| 254 | max_probs, max_idx = probs.max(dim=-1) |
| 255 | |
| 256 | # compute weight |
| 257 | mu = self.prob_max_mu_t |
| 258 | var = self.prob_max_var_t |
| 259 | mask = torch.exp(-((torch.clamp(max_probs - mu, max=0.0) ** 2) / (2 * var / 4))) |
| 260 | return max_probs.detach(), mask.detach() |
| 261 | |
| 262 | @torch.no_grad() |
| 263 | def distribution_alignment(self, probs): |