Improve Robustness on CPU machines

allenai · Oct 23, 2021 · b4c2e7f · b4c2e7f
1 parent 28bf44c
commit b4c2e7f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/vila/predictors.py b/src/vila/predictors.py
@@ -74,7 +74,7 @@ def predict(self, pdf_data) -> lp.Layout:
         return self.postprocess_model_outputs(pdf_data, model_inputs, model_predictions)
 
     def get_category_prediction(self, model_outputs):
-        predictions = model_outputs.logits.argmax(dim=-1).cpu().numpy()
+        predictions = model_outputs.logits.argmax(dim=-1).cpu().detach().numpy()
         return predictions
 
     def preprocess_pdf_data(self, pdf_data):