Fix CUDA memory issues

2021-03-30 09:02:41 +00:00
parent 5fb900bcb1
commit 9fc2c991f5
7 changed files with 72 additions and 60 deletions
--- a/lib/xlayers/super_norm.py
+++ b/lib/xlayers/super_norm.py
@@ -25,8 +25,8 @@ class SuperLayerNorm1D(SuperModule):
        self._eps = eps
        self._elementwise_affine = elementwise_affine
        if self._elementwise_affine:
-            self.weight = nn.Parameter(torch.Tensor(self.in_dim))
-            self.bias = nn.Parameter(torch.Tensor(self.in_dim))
+            self.register_parameter("weight", nn.Parameter(torch.Tensor(self.in_dim)))
+            self.register_parameter("bias", nn.Parameter(torch.Tensor(self.in_dim)))
        else:
            self.register_parameter("weight", None)
            self.register_parameter("bias", None)