Train. Args: mode: TODO.
(self, mode: bool = True)
| 442 | nn.init.zeros_(self.lora_B) |
| 443 | |
| 444 | def train(self, mode: bool = True): |
| 445 | """Train. |
| 446 | |
| 447 | Args: |
| 448 | mode: TODO. |
| 449 | """ |
| 450 | nn.Conv2d.train(self, mode) |
| 451 | if self.merge_weights and self.merged: |
| 452 | # Make sure that the weights are not merged |
| 453 | self.weight.data -= (self.lora_B @ self.lora_A).view(self.weight.shape) * self.scaling |
| 454 | self.merged = False |
| 455 | |
| 456 | def eval(self): |
| 457 | """Eval.""" |