mirror of
https://github.com/JDAI-CV/fast-reid.git
synced 2025-06-03 14:50:47 +08:00
bugfix for cuda tensor to numpy
Move cuda tensor to host memory in predictor which can be used for following process.
This commit is contained in:
parent
4f90197336
commit
62ad5e1a8b
@ -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):
|
||||
|
Loading…
x
Reference in New Issue
Block a user