Update mono/model/decode_heads/RAFTDepthNormalDPTDecoder5.py

mono/model/decode_heads/RAFTDepthNormalDPTDecoder5.py

@@ -792,8 +792,8 @@ class RAFTDepthNormalDPT5(nn.Module):
         self.relu = nn.ReLU(inplace=True)
     
     def get_bins(self, bins_num):
-        #depth_bins_vec = torch.linspace(math.log(self.min_val), math.log(self.max_val), bins_num, device="cuda")
-        depth_bins_vec = torch.linspace(math.log(self.min_val), math.log(self.max_val), bins_num, device="cpu")
+        depth_bins_vec = torch.linspace(math.log(self.min_val), math.log(self.max_val), bins_num, device="cuda")
+        #depth_bins_vec = torch.linspace(math.log(self.min_val), math.log(self.max_val), bins_num, device="cpu")
         depth_bins_vec = torch.exp(depth_bins_vec)
         return depth_bins_vec
     
@@ -848,7 +848,8 @@ class RAFTDepthNormalDPT5(nn.Module):
         return norm_normalize(torch.cat([normal_out, confidence], dim=1))
         #return norm_normalize(torch.cat([normal_out, confidence], dim=1).float())
     
-    def create_mesh_grid(self, height, width, batch, device="cpu", set_buffer=True):
+    #def create_mesh_grid(self, height, width, batch, device="cpu", set_buffer=True):
+    def create_mesh_grid(self, height, width, batch, device="cuda", set_buffer=True):
         y, x = torch.meshgrid([torch.arange(0, height, dtype=torch.float32, device=device),
                                torch.arange(0, width, dtype=torch.float32, device=device)], indexing='ij')
         meshgrid = torch.stack((x, y))