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))