mirrors
/
fast-reid
mirror of https://github.com/JDAI-CV/fast-reid.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

bugfix for cuda tensor to numpy 

Move cuda tensor to host memory in predictor which can be used for following process.
pull/525/head
liaoxingyu 2021-06-17 16:05:45 +08:00
parent 4f90197336
commit 62ad5e1a8b
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
fastreid/engine/defaults.py
View File

@ -151,7 +151,7 @@ class DefaultPredictor:
inputs = {"images": image.to(self.model.device)} inputs = {"images": image.to(self.model.device)}
with torch.no_grad(): # https://github.com/sphinx-doc/sphinx/issues/4258 with torch.no_grad(): # https://github.com/sphinx-doc/sphinx/issues/4258
predictions = self.model(inputs) predictions = self.model(inputs)
return predictions return predictions.cpu()
class DefaultTrainer(TrainerBase): class DefaultTrainer(TrainerBase):