def collate(self, batch, is_train=True):
if is_train:
packer = Packer([x['teller_tokens_in'] for x in batch])
return {
'packer':
packer,
'brw_teller_tokens_in':
packer.brw_from_list([x['teller_tokens_in']
for x in batch]).to(cuda_if_available),
'brw_teller_counts_in':
packer.brw_from_list([x['teller_counts_in']
for x in batch]).to(cuda_if_available),
'brw_teller_tokens_out':
packer.brw_from_list([x['teller_tokens_out']
for x in batch]).to(cuda_if_available),
'b_scene_tags':
torch.stack([x['scene_tags']
for x in batch]).to(cuda_if_available),
'b_scene_mask':
torch.stack([x['scene_mask']
for x in batch]).to(cuda_if_available),
'br_drawer_clipart_state':
packer.br_from_list([x['drawer_clipart_state']
for x in batch]).to(cuda_if_available),
}
else:
return {
'b_scene_tags':
torch.stack([x['scene_tags']
for x in batch]).to(cuda_if_available),
'b_scene_mask':
torch.stack([x['scene_mask']
for x in batch]).to(cuda_if_available),
}
