Rounding to nearest pixel value breaks almost all attacks about auto-attack HOT 3 CLOSED

import torch

class CustomRound(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        return torch.round(x)

    @staticmethod
    def backward(ctx, g):
        # send the gradient g straight-through on the backward pass.
        return g, None

class ModelWrapper(torch.nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model.eval()
        self.mean = torch.tensor([0.4914, 0.4822, 0.4465]).view(1, 3, 1, 1).to(device)
        self.std = torch.tensor([0.2470, 0.2435, 0.2616]).view(1, 3, 1, 1).to(device)
        self.round = CustomRound.apply
    
    def forward(self, x):
        x = x.clamp(0, 1)
        x = x * 255
        x = self.round(x)
        x = x / 255
        x = (x - self.mean) / self.std
        x = self.model(x)
        return x

model = ModelWrapper(resnet18(pretrained=True))
model = model.to(device).eval()