Update predict_from_net.py

justusschock · May 4, 2020 · 2501b68 · 2501b68
1 parent 260a759
commit 2501b68
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/shapenet/scripts/predict_from_net.py b/shapenet/scripts/predict_from_net.py
@@ -115,7 +115,7 @@ def process_sample(sample, img_size, net, device, crop=0.1):
         ).to(torch.float).unsqueeze(0).to(device)
 
         # obtain prediction
-        pred = net(img_tensor).cpu().numpy()["pred"][0]
+        pred = net(img_tensor)['pred'][0].cpu().numpy() 
 
         # remap to original image
         pred = pred * np.array([max_range / img_size, max_range / img_size])