From 07872d4cdb5532395e7c67fdf15477987fcc194d Mon Sep 17 00:00:00 2001 From: liaoxingyu Date: Thu, 20 Aug 2020 16:28:52 +0800 Subject: [PATCH] fix caffe export bug Summary: put tensor to gpu device --- tools/deploy/caffe_export.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/deploy/caffe_export.py b/tools/deploy/caffe_export.py index db89612..e6651b7 100644 --- a/tools/deploy/caffe_export.py +++ b/tools/deploy/caffe_export.py @@ -69,7 +69,7 @@ if __name__ == '__main__': model.eval() logger.info(model) - inputs = torch.randn(1, 3, cfg.INPUT.SIZE_TEST[0], cfg.INPUT.SIZE_TEST[1]) + inputs = torch.randn(1, 3, cfg.INPUT.SIZE_TEST[0], cfg.INPUT.SIZE_TEST[1]).to(torch.device(cfg.MODEL.DEVICE)) PathManager.mkdirs(args.output) pytorch_to_caffe.trans_net(model, inputs, args.name) pytorch_to_caffe.save_prototxt(f"{args.output}/{args.name}.prototxt")