diff --git a/neural_filters/neural_filter.py b/neural_filters/neural_filter.py
index a7550c23ab2770aca8b991920cfd3c71b2e67d9c..4ee01679b293be73f5bf902c9ed810e4c124b247 100644
--- a/neural_filters/neural_filter.py
+++ b/neural_filters/neural_filter.py
@@ -34,7 +34,7 @@ class NeuralFilter(torch.nn.Module):
 
         self.hidden_size = hidden_size
 
-        self.bias_forget = Parameter(torch.Tensor(hidden_size))
+        self.register_parameter('bias_forget', torch.Tensor(hidden_size))
 
         self.reset_parameters()
 
diff --git a/neural_filters/neural_filter_2CC.py b/neural_filters/neural_filter_2CC.py
index 6e2e294677b8b9219641d20b491107809cd38e05..03f49ab36db1201614a3b979c5be63c708e14aa2 100644
--- a/neural_filters/neural_filter_2CC.py
+++ b/neural_filters/neural_filter_2CC.py
@@ -35,8 +35,8 @@ class NeuralFilter2CC(torch.nn.Module):
 
         self.hidden_size = hidden_size
 
-        self.bias_theta = Parameter(torch.Tensor(hidden_size))
-        self.bias_modulus = Parameter(torch.Tensor(hidden_size))
+        self.register_parameter('bias_theta', torch.Tensor(hidden_size))
+        self.register_parameter('bias_modulus', torch.Tensor(hidden_size))
 
         self.reset_parameters()
 
diff --git a/setup.py b/setup.py
index d37047e05931bfd1000659a05893d674cf56a4da..3022b39c96a1f889d7f46555bea9f24b0854e0a4 100644
--- a/setup.py
+++ b/setup.py
@@ -2,7 +2,7 @@ from setuptools import setup
 
 setup(
     name='neural-filters',
-    version='0.1',
+    version='0.2',
     description='Linear filters for neural networks in pyTorch',
     author='François Marelli (Idiap research institute)',
     author_email='francois.marelli@idiap.ch',