Merge branch 'master' of https://github.com/zju3dv/pvnet

zju3dv · May 22, 2019 · bac3608 · bac3608
2 parents 3392b5e + 4b4a444
commit bac3608
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tools/demo.py b/tools/demo.py
@@ -104,16 +104,16 @@ def read_data():
 
 
 def visualize_mask(mask):
- plt.imshow(mask[0])
+ plt.imshow(mask[0].cpu())
  plt.show()
 
 
 def visualize_vertex(vertex, vertex_weights):
  vertex = vertex * vertex_weights
  for i in range(9):
  _, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 8))
- ax1.imshow(vertex[0, 2*i])
- ax2.imshow(vertex[0, 2*i+1])
+ ax1.imshow(vertex[0, 2*i].cpu())
+ ax2.imshow(vertex[0, 2*i+1].cpu())
  plt.show()