diff --git a/dac/model/discriminator.py b/dac/model/discriminator.py
index 09c79d1..329a420 100644
--- a/dac/model/discriminator.py
+++ b/dac/model/discriminator.py
@@ -5,7 +5,7 @@
 from audiotools import ml
 from audiotools import STFTParams
 from einops import rearrange
-from torch.nn.utils import weight_norm
+from torch.nn.utils.parametrizations import weight_norm
 
 
 def WNConv1d(*args, **kwargs):
diff --git a/dac/nn/layers.py b/dac/nn/layers.py
index 44fbc29..81ec7ef 100644
--- a/dac/nn/layers.py
+++ b/dac/nn/layers.py
@@ -1,9 +1,7 @@
 import numpy as np
 import torch
 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):
diff --git a/dac/nn/quantize.py b/dac/nn/quantize.py
index b17ff4a..cc42a6a 100644
--- a/dac/nn/quantize.py
+++ b/dac/nn/quantize.py
@@ -5,7 +5,6 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from einops import rearrange
-from torch.nn.utils import weight_norm
 
 from dac.nn.layers import WNConv1d
 
diff --git a/requirements.txt b/requirements.txt
index 2919657..b0e24aa 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,7 +2,7 @@ argbind>=0.3.7
 descript-audiotools>=0.7.2
 einops
 numpy
-torch
+torch>= 2.1.0
 torchaudio
 tqdm
 tensorboard