move model to gpu in tools/visactmap

2019-11-27 17:01:32 +00:00 · 2019-11-27 17:01:32 +00:00 · 67f56ca361
parent bf79ba96be
commit 67f56ca361
1 changed files with 3 additions and 0 deletions
--- a/tools/visualize_actmap.py
+++ b/tools/visualize_actmap.py
@ -133,6 +133,9 @@ def main():
        num_classes=datamanager.num_train_pids,
        use_gpu=use_gpu
    )
+
+    if use_gpu:
+        model = model.cuda()
    
    if args.weights and check_isfile(args.weights):
        load_pretrained_weights(model, args.weights)