update

models/DiMeR/models/DiMeR.py

@@ -83,6 +83,7 @@ class DiMeR(nn.Module):
             samples_per_ray=rendering_samples_per_ray,
             use_deformation_weight=use_deformation_weight,
         )
+        
     @spaces.GPU
     def init_flexicubes_geometry(self, device, fovy=50.0):
         camera = PerspectiveCamera(fovy=fovy, device=device)

models/DiMeR/models/geometry/render/neural_render.py

@@ -7,6 +7,7 @@
 # license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
 
 import torch
+import spaces
 import torch.nn.functional as F
 import nvdiffrast.torch as dr
 from . import Renderer
@@ -146,6 +147,7 @@ class NeuralRender(Renderer):
     #  - Single light
     #  - Single material
     # ==============================================================================================
+    @spaces.GPU
     def render_layer(
             self,
             rast,
@@ -189,6 +191,7 @@ class NeuralRender(Renderer):
 
         return gb_pos, gb_normal
 
+    @spaces.GPU
     def render_mesh(
             self,
             mesh_v_pos_bxnx3,
@@ -245,6 +248,7 @@ class NeuralRender(Renderer):
         # normal = torch.lerp(torch.zeros_like(normal), (normal + 1.0) / 2.0, hard_mask.float())      # black background
         return ori_mesh_feature, antialias_mask, hard_mask, rast, v_pos_clip, mask_pyramid, depth, normal, gb_normal
     
+    @spaces.GPU
     def render_mesh_light(
             self,
             mesh_v_pos_bxnx3,