(self,in_features,out_features,gain = 2)
| 48 | |
| 49 | class WSLinear(torch.nn.Module): |
| 50 | def __init__(self,in_features,out_features,gain = 2): |
| 51 | super(WSLinear, self).__init__() |
| 52 | self.linear = torch.nn.Linear(in_features,out_features) |
| 53 | self.scale = (gain / in_features)**0.5 |
| 54 | self.bias = self.linear.bias |
| 55 | self.linear.bias = None |
| 56 | |
| 57 | #initialize linear layer |
| 58 | torch.nn.init.normal_(self.linear.weight) |
| 59 | torch.nn.init.zeros_(self.bias) |
| 60 | |
| 61 | def forward(self,x): |
| 62 | out = self.linear(x * self.scale) + self.bias |