| 312 | """size of each output sample.""" |
| 313 | |
| 314 | def __init__( |
| 315 | self, |
| 316 | weight: Union[torch.Tensor, nn.Parameter], |
| 317 | bias: Union[None, torch.Tensor, nn.Parameter], |
| 318 | ): |
| 319 | super().__init__() |
| 320 | |
| 321 | if weight.ndim != 2: |
| 322 | raise ValueError( |
| 323 | f"weight parameter has {weight.ndim} dimensions, should have 2" |
| 324 | ) |
| 325 | self.out_features, self.in_features = weight.shape |
| 326 | |
| 327 | self.register_buffer("weight", weight) |
| 328 | self.register_buffer("bias", bias) |
| 329 | |
| 330 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 331 | return F.linear(x, self.weight, self.bias) |