xrg commited on
Commit
5708415
·
1 Parent(s): 72f6244
Files changed (2) hide show
  1. core/models.py +5 -4
  2. core/tensorBase.py +2 -2
core/models.py CHANGED
@@ -48,10 +48,11 @@ class LTRFM_NeRF(nn.Module):
48
  app_n_comp=opt.app_n_comp,
49
  density_n_comp=opt.density_n_comp,
50
  )
51
-
52
- aabb = torch.tensor([[-1, -1, -1], [1, 1, 1]]).to(device)
53
- grid_size = torch.tensor([opt.splat_size, opt.splat_size, opt.splat_size]).to(device)
54
- near_far =torch.tensor([opt.znear, opt.zfar]).to(device)
 
55
 
56
  print(device)
57
  print('aabb')
 
48
  app_n_comp=opt.app_n_comp,
49
  density_n_comp=opt.density_n_comp,
50
  )
51
+ print(aabb.device)
52
+ aabb = torch.tensor([[-1, -1, -1], [1, 1, 1]]).cuda()
53
+ print(aabb.device)
54
+ grid_size = torch.tensor([opt.splat_size, opt.splat_size, opt.splat_size]).cuda()
55
+ near_far =torch.tensor([opt.znear, opt.zfar]).cuda()
56
 
57
  print(device)
58
  print('aabb')
core/tensorBase.py CHANGED
@@ -243,8 +243,8 @@ class TensorBase(torch.nn.Module):
243
  print(self.aabbSize.device)
244
  self.invaabbSize = 2.0/self.aabbSize
245
  print(self.invaabbSize.device)
246
- self.invaabbSize = self.invaabbSize.to(self.aabb.device)
247
- print(self.invaabbSize.device)
248
  self.gridSize= gridSize.float()
249
  self.units=self.aabbSize / (self.gridSize-1)
250
  self.stepSize=torch.mean(self.units)*self.step_ratio # TBD step_ratio? why so small 0.5
 
243
  print(self.aabbSize.device)
244
  self.invaabbSize = 2.0/self.aabbSize
245
  print(self.invaabbSize.device)
246
+ # self.invaabbSize = self.invaabbSize.to(self.aabb.device)
247
+ # print(self.invaabbSize.device)
248
  self.gridSize= gridSize.float()
249
  self.units=self.aabbSize / (self.gridSize-1)
250
  self.stepSize=torch.mean(self.units)*self.step_ratio # TBD step_ratio? why so small 0.5