bugfix for cuda tensor to numpy

Move cuda tensor to host memory in predictor which can be used for following process.
2021-06-17 16:05:45 +08:00 · 2021-06-17 16:05:45 +08:00 · 62ad5e1a8b
parent 4f90197336
commit 62ad5e1a8b
1 changed files with 1 additions and 1 deletions
--- a/fastreid/engine/defaults.py
+++ b/fastreid/engine/defaults.py
@ -151,7 +151,7 @@ class DefaultPredictor:
        inputs = {"images": image.to(self.model.device)}
        with torch.no_grad():  # https://github.com/sphinx-doc/sphinx/issues/4258
            predictions = self.model(inputs)
-        return predictions
+        return predictions.cpu()


 class DefaultTrainer(TrainerBase):