s444415/donut
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

working streaming training

Browse Source
This commit is contained in:
michal.kozlowski 2023-03-17 15:18:40 +01:00
parent d9cbe1f9a8
commit 34316df6f5
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
utils/donut_model_pl_stream.py
View File

@ -29,7 +29,7 @@ class DonutModelPLModuleStream(pl.LightningModule):
pixel_values = batch['pixel_values'] pixel_values = batch['pixel_values']
labels = batch['labels'] labels = batch['labels']
answers = batch['target_sequence'] answers = batch['target_sequence'][0]
batch_size = pixel_values.shape[0] batch_size = pixel_values.shape[0]
# we feed the prompt to the model # we feed the prompt to the model
decoder_input_ids = torch.full((batch_size, 1), self.model.config.decoder_start_token_id, device=self.device) decoder_input_ids = torch.full((batch_size, 1), self.model.config.decoder_start_token_id, device=self.device)