def augment_batch(self, batch: Tensor) -> Tensor:
        """Generate anomalous masks for a batch of images.

        Args:
            batch (Tensor): Input batch of images of shape
                ``(batch_size, channels, height, width)``.

        Returns:
            Tensor: Batch of binary masks of shape
                ``(batch_size, 1, height, width)`` where ``1`` indicates
                anomalous regions.

        Example:
            >>> generator = DsrAnomalyGenerator()
            >>> batch = torch.randn(8, 3, 256, 256)
            >>> masks = generator.augment_batch(batch)
            >>> assert masks.shape == (8, 1, 256, 256)
            >>> assert torch.all((masks >= 0) & (masks <= 1))
        """
        batch_size, _, height, width = batch.shape

        # Collect perturbations
        masks_list: list[Tensor] = []
        for _ in range(batch_size):
            if torch.rand(1) > self.p_anomalous:  # include normal samples
                masks_list.append(torch.zeros((1, height, width), device=batch.device))
            else:
                mask = self.generate_anomaly(height, width)
                masks_list.append(mask)

        return torch.stack(masks_list).to(batch.device)
Generate anomalous masks for a batch of images.

        Args:
            batch (Tensor): Input batch of images of shape
                ``(batch_size, channels, height, width)``.

        Returns:
            Tensor: Batch of binary masks of shape
                ``(batch_size, 1, height, width)`` where ``1`` indicates
                anomalous regions.

        Example:
            >>> generator = DsrAnomalyGenerator()
            >>> batch = torch.randn(8, 3, 256, 256)
            >>> masks = generator.augment_batch(batch)
            >>> assert masks.shape == (8, 1, 256, 256)
            >>> assert torch.all((masks >= 0) & (masks <= 1))
        
src/anomalib/models/image/dsr/anomaly_generator.py