In order to work with newer versions of pytorch

Converting CUDA tensor to numpy was only possible after copying tensors to host memory. Implicit method was removed in #9934.
zju3dv · May 14, 2019 · 7bc57c9 · 7bc57c9
1 parent fbde7ba
commit 7bc57c9
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()