Hi,
From your code, there is no “infinity”, you do this norm by doing torch.max(torch.abs(w_tmp))**2. If you were doing w_tmp.norm(2) then one is doing one operation while the other is doing 3.
What is the memory usage for the corresponding code that uses the l2 regularization?
If you could provide a small code snippet that reproduce the issue, that could be helpful.