| 130 | """ |
| 131 | |
| 132 | def apply(self, data: torch.Tensor): |
| 133 | weights, perm, coords = self._params |
| 134 | nsamples, _, *dims = data.shape |
| 135 | if len(weights) != nsamples: |
| 136 | raise ValueError(f"Expected batch of size: {len(weights)}, but got {nsamples}") |
| 137 | |
| 138 | mask = torch.ones_like(data) |
| 139 | for s, weight in enumerate(weights): |
| 140 | lengths = [d * sqrt(1 - weight) for d in dims] |
| 141 | idx = [slice(None)] + [slice(c, min(ceil(c + ln), d)) for c, ln, d in zip(coords, lengths, dims)] |
| 142 | mask[s][idx] = 0 |
| 143 | |
| 144 | return mask * data + (1 - mask) * data[perm, ...] |
| 145 | |
| 146 | def apply_on_labels(self, labels: torch.Tensor): |
| 147 | weights, perm, _ = self._params |