new

models/ISOMER/scripts/project_mesh.py

@@ -13,6 +13,7 @@ from pytorch3d.renderer import (
     FoVOrthographicCameras,
 )
 from pytorch3d.renderer import MeshRasterizer
+import spaces
 
 def get_camera(world_to_cam, fov_in_degrees=60, focal_length=1 / (2**0.5), cam_type='fov'):
     # pytorch3d expects transforms as row-vectors, so flip rotation: https://github.com/facebookresearch/pytorch3d/issues/1183
@@ -68,6 +69,7 @@ def _warmup(glctx, device=None):
     dr.rasterize(glctx, pos, tri, resolution=[256, 256])
 
 class Pix2FacesRenderer:
+    @spaces.GPU
     def __init__(self, device="cuda"):
         # self._glctx = dr.RasterizeGLContext(output_db=False, device=device)
         self._glctx = dr.RasterizeCudaContext(device=device)