From 67f56ca3613fb17ddb3c4b29d7b43abbe0623847 Mon Sep 17 00:00:00 2001
From: KaiyangZhou <k.zhou@surrey.ac.uk>
Date: Wed, 27 Nov 2019 17:01:32 +0000
Subject: [PATCH] move model to gpu in tools/visactmap

---
 tools/visualize_actmap.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/tools/visualize_actmap.py b/tools/visualize_actmap.py
index 278d25e..e148050 100644
--- 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)