Thanks for the tip. I wrote this code:
def normalize_channels( _x, inplace=False):
        tmp = torch.flatten(_x, start_dim=2)
        _min = tmp.min(dim=-1)[0]
        _max = tmp.max(dim=-1)[0]
        del tmp
        if inplace:
            return _x.add_(_min.unsqueeze(-1).unsqueeze(-1), alpha=-1).mul_(
                1 / _max.add(_min, alpha=-1).unsqueeze(-1).unsqueeze(-1))
        else:
            return _x.add(_min.unsqueeze(-1).unsqueeze(-1), alpha=-1).mul(
                1 / _max.add(_min, alpha=-1).unsqueeze(-1).unsqueeze(-1))

And it performs better:




ITERATIONS
OLD (s)
NEW (s)
OLD-NEW (s)




0
2,42E-01
7,80E-03
2,34E-01


1
2,38E-01
8,23E-03
2,30E-01


2
2,34E-01
8,17E-03
2,26E-01


3
2,26E-01
8,22E-03
2,18E-01


4
2,24E-01
8,21E-03
2,16E-01


5
2,54E-01
8,19E-03
2,46E-01


mean
2,36E-01
8,14E-03
2,28E-01