hungdungn47 commited on
Commit
d91835f
·
1 Parent(s): 0e04b12

reduce num worker

Browse files
Files changed (1) hide show
  1. infer_concat.py +1 -1
infer_concat.py CHANGED
@@ -104,6 +104,6 @@ def infer_2_hier(model, data_loader, device, tokenizer):
104
 
105
  def vit5_infer(data):
106
  dataset = Dataset4Summarization(data, tokenizer)
107
- data_loader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=2)
108
  result = infer_2_hier(model, data_loader, device, tokenizer)
109
  return result
 
104
 
105
  def vit5_infer(data):
106
  dataset = Dataset4Summarization(data, tokenizer)
107
+ data_loader = torch.utils.data.DataLoader(dataset, batch_size=1, num_workers=1)
108
  result = infer_2_hier(model, data_loader, device, tokenizer)
109
  return result