Update models/ISOMER/mesh_reconstruction/render.py

models/ISOMER/mesh_reconstruction/render.py

@@ -2,6 +2,7 @@
 import nvdiffrast.torch as dr
 import torch
 from typing import Tuple
+import spaces
 
 def _warmup(glctx, device=None):
     device = 'cuda' if device is None else device
@@ -15,6 +16,7 @@ def _warmup(glctx, device=None):
     dr.rasterize(glctx, pos, tri, resolution=[256, 256])
 
 # glctx = dr.RasterizeGLContext(output_db=False, device="cuda")
+@spaces.GPU
 glctx = dr.RasterizeCudaContext(device="cuda")
 
 class NormalsRenderer: