iicd/autodl-projects
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix CUDA memory issues

Browse Source
This commit is contained in:
D-X-Y
2021-03-30 09:02:41 +00:00
parent 5fb900bcb1
commit 9fc2c991f5
7 changed files with 72 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
lib/xlayers/super_linear.py
View File

@@ -30,11 +30,14 @@ class SuperLinear(SuperModule):
self._out_features = out_features
self._bias = bias
# weights to be optimized
self._super_weight = torch.nn.Parameter(
torch.Tensor(self.out_features, self.in_features)
self.register_parameter(
"_super_weight",
torch.nn.Parameter(torch.Tensor(self.out_features, self.in_features)),
)
if self.bias:
self._super_bias = torch.nn.Parameter(torch.Tensor(self.out_features))
self.register_parameter(
"_super_bias", torch.nn.Parameter(torch.Tensor(self.out_features))
)
else:
self.register_parameter("_super_bias", None)
self.reset_parameters()

4
lib/xlayers/super_norm.py
View File

@@ -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)