| 179 | """ |
| 180 | |
| 181 | def apply(self, data: torch.Tensor): |
| 182 | weights, _, coords = self._params |
| 183 | nsamples, _, *dims = data.shape |
| 184 | if len(weights) != nsamples: |
| 185 | raise ValueError(f"Expected batch of size: {len(weights)}, but got {nsamples}") |
| 186 | |
| 187 | mask = torch.ones_like(data) |
| 188 | for s, weight in enumerate(weights): |
| 189 | lengths = [d * sqrt(1 - weight) for d in dims] |
| 190 | idx = [slice(None)] + [slice(c, min(ceil(c + ln), d)) for c, ln, d in zip(coords, lengths, dims)] |
| 191 | mask[s][idx] = 0 |
| 192 | |
| 193 | return mask * data |
| 194 | |
| 195 | def __call__(self, data: torch.Tensor, randomize=True): |
| 196 | data_t = convert_to_tensor(data, track_meta=get_track_meta()) |