| 23 | setattr(obj, parts[-1], value) |
| 24 | |
| 25 | class DataProvider: |
| 26 | def __init__(self): |
| 27 | self.batch = None |
| 28 | |
| 29 | def set_batch(self, batch): |
| 30 | if self.batch is not None: |
| 31 | if isinstance(self.batch, torch.Tensor): |
| 32 | assert self.batch.shape[1:] == batch.shape[1:], "Check: shapes probably should not change during training" |
| 33 | |
| 34 | self.batch = batch |
| 35 | |
| 36 | def get_batch(self, x=None): |
| 37 | assert self.batch is not None, "Error: need to set a batch first" |
| 38 | |
| 39 | if x is None or isinstance(self.batch, torch.Tensor): |
| 40 | return self.batch |
| 41 | |
| 42 | # batch is a list; select the corresponding element based on x |
| 43 | size = x.shape[2] |
| 44 | for i in range(len(self.batch)): |
| 45 | if self.batch[i].shape[2] == size: |
| 46 | return self.batch[i] |
| 47 | |
| 48 | raise ValueError("Error: no matching batch found") |
| 49 | |
| 50 | def reset(self): |
| 51 | self.batch = None |
| 52 | |
| 53 | class LoraLinear(torch.nn.Module): |
| 54 | def __init__( |