daidedou commited on
Commit
6cc6bd1
·
1 Parent(s): c5cb8fa

real cpu need

Browse files
Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -267,6 +267,8 @@ if __name__ == "__main__":
267
  parser.add_argument('--share', action="store_true")
268
  args = parser.parse_args()
269
  cfg = OmegaConf.load(args.config)
270
- matcher = zero_shot.Matcher(cfg)
 
 
271
  datadicts = None
272
  demo.launch(share=args.share)
 
267
  parser.add_argument('--share', action="store_true")
268
  args = parser.parse_args()
269
  cfg = OmegaConf.load(args.config)
270
+ device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
271
+ print(f"Using device: {device}")
272
+ matcher = zero_shot.Matcher(cfg, device)
273
  datadicts = None
274
  demo.launch(share=args.share)