net architecture fixed

BaiCHao12 · Oct 13, 2019 · aa0ef2f · aa0ef2f
1 parent 3f32895
commit aa0ef2f
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 3 deletions.
diff --git a/eval.py b/eval.py
@@ -59,7 +59,7 @@
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-net.load_state_dict(torch.load(checkpoint_fname)['state_dict'])
+net.load_state_dict(torch.load(checkpoint_fname, map_location=device)['state_dict'])
 net = nn.DataParallel(net)
 net.eval()
 net = net.to(device)

diff --git a/model/net.py b/model/net.py
@@ -64,8 +64,6 @@ def forward(self, x1, x2):
         # do correlation
         corr1 = self.corr(feat_top_pyr_trg, feat_top_pyr_src)
         corr1 = self.l2norm(F.relu(corr1))
-        print(corr1.shape)
-        sys.exit()
 
         b, c, h, w = corr1.size()
         init_map = torch.FloatTensor(b, 2, h, w).zero_().to(self.device)