Train. Args: mode: TODO.
(self, mode: bool = True)
| 81 | nn.init.normal_(self.lora_B) |
| 82 | |
| 83 | def train(self, mode: bool = True): |
| 84 | """Train. |
| 85 | |
| 86 | Args: |
| 87 | mode: TODO. |
| 88 | """ |
| 89 | nn.Embedding.train(self, mode) |
| 90 | if self.merge_weights and self.merged: |
| 91 | # Make sure that the weights are not merged |
| 92 | if self.r > 0: |
| 93 | self.weight.data -= (self.lora_B @ self.lora_A).T * self.scaling |
| 94 | self.merged = False |
| 95 | |
| 96 | def eval(self): |
| 97 | """Eval.""" |