Making DAC decoder torch.compileable

The original torch.nn.utils.weight_norm() function is deprecated and it prevents the DAC decoder model from being torch.compiled. The newer torch.nn.utils.parametrizations.weight_norm() works the same way as the old one and plus it is compatible with torch.compile (with CUDA graphs).

Author: Artyom17

dac/nn/layers.py

@@ -3,7 +3,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
-from torch.nn.utils import weight_norm
+from torch.nn.utils.parametrizations import weight_norm
 
 
 def WNConv1d(*args, **kwargs):