fix a bug in trainers/pvnet.py

willer94 · Mar 9, 2020 · 2ec747f · 2ec747f
1 parent b9ef443
commit 2ec747f
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 2 deletions.
diff --git a/lib/datasets/linemod/pvnet.py b/lib/datasets/linemod/pvnet.py
@@ -51,7 +51,7 @@ def __getitem__(self, index_tuple):
             inp, kpt_2d, mask = self._transforms(inp, kpt_2d, mask)
 
         vertex = pvnet_data_utils.compute_vertex(mask, kpt_2d).transpose(2, 0, 1)
-        ret = {'inp': inp, 'mask': mask.astype(np.uint8), 'vertex': vertex, 'img_id': img_id, 'meta': ''}
+        ret = {'inp': inp, 'mask': mask.astype(np.uint8), 'vertex': vertex, 'img_id': img_id, 'meta': {}}
         # visualize_utils.visualize_linemod_ann(torch.tensor(inp), kpt_2d, mask, True)
 
         return ret

diff --git a/lib/train/trainers/pvnet.py b/lib/train/trainers/pvnet.py
@@ -19,6 +19,7 @@ def forward(self, batch):
         loss = 0
 
         if 'pose_test' in batch['meta'].keys():
+            loss = torch.tensor(0).to(batch['inp'].device)
             return output, loss, {}, {}
 
         weight = batch['mask'][:, None].float()

diff --git a/lib/train/trainers/trainer.py b/lib/train/trainers/trainer.py
@@ -79,7 +79,7 @@ def val(self, epoch, data_loader, evaluator=None, recorder=None):
                     batch[k] = batch[k].cuda()
 
             with torch.no_grad():
-                output, loss, loss_stats, image_stats = self.network(batch)
+                output, loss, loss_stats, image_stats = self.network.module(batch)
                 if evaluator is not None:
                     evaluator.evaluate(output, batch)