(batch)
| 252 | |
| 253 | |
| 254 | def generate_batch(batch): |
| 255 | title, files, author, category = zip(*batch) |
| 256 | title = list(title) |
| 257 | files = list(files) |
| 258 | author = list(author) |
| 259 | category = list(category) |
| 260 | targets = torch.tensor([common.categories.index(cat) for cat in category]).to( |
| 261 | device |
| 262 | ) |
| 263 | return CommitClassifierInputs(title, files, author), targets |
| 264 | |
| 265 | |
| 266 | def train_step(batch, model, optimizer, loss): |