#147 related

nicolay-r · Jul 13, 2024 · e18e744 · e18e744
1 parent 883ef7e
commit e18e744
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/arelight/pipelines/items/inference_bert_opennre.py b/arelight/pipelines/items/inference_bert_opennre.py
@@ -137,7 +137,7 @@ def iter_results(parallel_model, eval_loader, data_ids):
                     # Save result
                     batch_size = pred.size(0)
                     for i in range(batch_size):
-                        yield data_ids[l_ind], pred[i].item()
+                        yield data_ids[l_ind], int(pred[i].item())
                         l_ind += 1
 
     def __iter_predict_result(self, samples_filepath, batch_size):