iicd/xautodl
1
0
Fork 0
Code Issues Pull Requests Actions 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()