Commit
·
d91835f
1
Parent(s):
0e04b12
reduce num worker
Browse files- 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=
|
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
|