Repository: Project-MONAI/MONAI
File: monai/transforms/regularization/array.py (Python)

Code before:
        if len(weights) != nsamples:
            raise ValueError(f"Expected batch of size: {len(weights)}, but got {nsamples}")

        mask = torch.ones_like(data)
        for s, weight in enumerate(weights):
            lengths = [d * sqrt(1 - weight) for d in dims]
            idx = [slice(None)] + [slice(c, min(ceil(c + ln), d)) for c, ln, d in zip(coords, lengths, dims)]
            mask[s][idx] = 0

        return mask * data + (1 - mask) * data[perm, ...]

    def apply_on_labels(self, labels: torch.Tensor):
        weights, perm, _ = self._params
        nsamples, *dims = labels.shape
        if len(weights) != nsamples:
            raise ValueError(f"Expected batch of size: {len(weights)}, but got {nsamples}")

        mixweight = weights[(Ellipsis,) + (None,) * len(dims)]
        return mixweight * labels + (1 - mixweight) * labels[perm, ...]

    def __call__(self, data: torch.Tensor, labels: torch.Tensor | None = None, randomize=True):
        data = convert_to_tensor(data, track_meta=get_track_meta())
        data_t = convert_to_tensor(data, track_meta=False)
        if labels is not None:
            labels_t = convert_to_tensor(labels, track_meta=get_track_meta())
            labels_t = convert_to_tensor(labels, track_meta=False)
        if randomize:
            self.randomize(data)
        augmented = convert_to_dst_type(self.apply(data_t), dst=data)[0]
        if labels is not None:
            augmented_label = convert_to_dst_type(self.apply(labels_t), dst=labels)[0]
        return (augmented, augmented_label) if labels is not None else augmented


class CutOut(Mixer):
    """Cutout as described in the paper:
    Terrance DeVries, Graham W. Taylor.
    Improved Regularization of Convolutional Neural Networks with Cutout,
    arXiv:1708.04552

    Class derived from :py:class:`monai.transforms.Mixer`. See corresponding
    documentation for details on the constructor parameters. Here, alpha not only determines
    the mixing weight but also the size of the random rectangles being cut put.
    Please refer to the paper for details.
    """

    def apply(self, data: torch.Tensor):
        weights, _, coords = self._params
        nsamples, _, *dims = data.shape
        if len(weights) != nsamples:
            raise ValueError(f"Expected batch of size: {len(weights)}, but got {nsamples}")

Review comment (suggestion):
```suggestion
        data_t = convert_to_tensor(data, track_meta=get_track_meta())
        if labels is not None:
            labels_t = convert_to_tensor(labels, track_meta=get_track_meta())
```

Code after:

        mask = torch.ones_like(data)
        for s, weight in enumerate(weights):
            lengths = [d * sqrt(1 - weight) for d in dims]
            idx = [slice(None)] + [slice(c, min(ceil(c + ln), d)) for c, ln, d in zip(coords, lengths, dims)]
            mask[s][idx] = 0

        return mask * data + (1 - mask) * data[perm, ...]

    def apply_on_labels(self, labels: torch.Tensor):
        weights, perm, _ = self._params
        nsamples, *dims = labels.shape
        if len(weights) != nsamples:
            raise ValueError(f"Expected batch of size: {len(weights)}, but got {nsamples}")

        mixweight = weights[(Ellipsis,) + (None,) * len(dims)]
        return mixweight * labels + (1 - mixweight) * labels[perm, ...]

    def __call__(self, data: torch.Tensor, labels: torch.Tensor | None = None, randomize=True):
        data_t = convert_to_tensor(data, track_meta=get_track_meta())
        if labels is not None:
            labels_t = convert_to_tensor(labels, track_meta=get_track_meta())
        if randomize:
            self.randomize(data)
        augmented = convert_to_dst_type(self.apply(data_t), dst=data)[0]
        if labels is not None:
            augmented_label = convert_to_dst_type(self.apply(labels_t), dst=labels)[0]
        return (augmented, augmented_label) if labels is not None else augmented


class CutOut(Mixer):
    """Cutout as described in the paper:
    Terrance DeVries, Graham W. Taylor.
    Improved Regularization of Convolutional Neural Networks with Cutout,
    arXiv:1708.04552

    Class derived from :py:class:`monai.transforms.Mixer`. See corresponding
    documentation for details on the constructor parameters. Here, alpha not only determines
    the mixing weight but also the size of the random rectangles being cut put.
    Please refer to the paper for details.
    """

    def apply(self, data: torch.Tensor):
        weights, _, coords = self._params
        nsamples, _, *dims = data.shape
        if len(weights) != nsamples:
            raise ValueError(f"Expected batch of size: {len(weights)}, but got {nsamples}")

        mask = torch.ones_like(data)
        for s, weight in enumerate(weights):
            lengths = [d * sqrt(1 - weight) for d in dims]