Hi All! I am trying to write a new loss function, which takes infinity norm of the weights.
I am trying some thing on this line.
def l2_regu(mdl):
        l2_reg = None
        for W in mdl.parameters():
                if W.ndimension() < 2:
                        continue
                else:
                        w_tmp = W
                        if l2_reg is None:
                                l2_reg = (torch.max(torch.abs(w_tmp)))**2
                        else:   
                                l2_reg = l2_reg + (torch.max(torch.abs(w_tmp)))**2
      
        return l2_reg

Running this throws an error  saying:
File "train.py", line 208, in train
    oloss =  l2_regu(model)
  File "train.py", line 38, in l2_reg_ortho
    l2_reg = l2_reg + (torch.max(torch.abs(w_tmp)))**2
RuntimeError: cuda runtime error (2) : out of memory at /pytorch/torch/lib/THC/generic/THCStorage.cu:58
(torch3) bansa01@vita2:~/pytorch_wideres/tmp_inf/WideResNet-pytorch$

Every thing works good if I remove/replace with l2 loss, but with infinity norm it throws the above error.
Regards,
Nitin