Solved by MalumaDev in post #3 
              
              
                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=…